[Paper Review] 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 Science54 citations

TL;DR

GemNet-OC delivers state-of-the-art results on OC20 while being ~10x faster to train, and the paper analyzes how dataset size, diversity, and domain shift affect GNN design, introducing OC-2M as a representative subset for efficient development.

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.

Motivation & Objective

• Investigate whether GNN improvements on small datasets transfer to large, diverse molecular data

Proposed method

• Develop GemNet-OC baseline based on OC20 with multi-level interaction hierarchy and edge/atom embeddings
• Replace fixed distance cutoffs with a fixed number of nearest neighbors for graph construction
• Simplify and optimize basis functions to reduce computational cost of radial/angular terms
• Introduce a hierarchical interaction scheme including atom-atom, edge-atom, and edge-edge interactions with a long-range atom-level pathway
• Output and concatenate embeddings from interaction blocks to improve final energy/force predictions
• Provide open-source code and pretrained weights to enable reproducible development on OC20 and OC-2M

Experimental results

Research questions

• RQ1 How do four dataset properties (chemical diversity, system size, dataset size, and domain shift) affect GNN design decisions?
• RQ2 Can a model trained on OC20 achieve state-of-the-art performance with reduced training cost?
• RQ3 Is OC-2M a reliable subset that correlates with OC20 trends for faster development?
• RQ4 Do model component effects vary across small and large datasets, and across OC20 subsets?
• RQ5 What architectural and training strategies best enable fast, scalable learning for large molecular datasets?

Key findings

• GemNet-OC achieves state-of-the-art results on OC20 tasks and is ~10x faster to train than prior large models
• GemNet-OC outperforms previous models on OC-2M and OC20 training data, while using substantially less training data
• There is a notable discrepancy in how model components perform across small vs. large, diverse datasets
• OC-2M correlates well with full OC20 results, enabling faster, representative model development
• An interaction hierarchy with long-range atom-level pathways and optimized basis functions enables efficient handling of large systems