Source-linked AI summary
GemNet-OC: Developing Graph Neural Networks for Large and Diverse Molecular Simulation Datasets
Johannes Gasteiger, Muhammed Shuaibi, Anuroop Sriram, Stephan Günnemann, Zachary Ulissi, C. Lawrence Zitnick, Abhishek Das
TL;DR
Small, narrow datasets dominate GNN benchmarking, leaving open whether their model improvements transfer to larger and more diverse molecular systems. The paper develops GemNet-OC on OC20 and compares model choices across datasets. GemNet-OC improves on the prior OC20 state of the art, while OC-2M provides a cheaper proxy for model development.
Problem
It is unclear whether GNN improvements demonstrated on small, narrow molecular datasets generalize to larger and more diverse systems.
Source: §1 · p. 2Method
The paper develops GemNet-OC on OC20 and compares model components and hyperparameters across datasets and six OC20 subsets.
Source: §1 · p. 2Results
GemNet-OC outperforms the previous state of the art on OC20 by 16%, while model-component and hyperparameter effects can differ or reverse across datasets.
Source: §1 · p. 2, §7 · p. 12Takeaways & Limitations
OC-2M's model trends correlate with OC20, providing a lower-cost dataset for model development and baseline comparisons.
Source: §1 · p. 2, §7 · p. 12
Abstract
from arXiv · showhide
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.
1 Introduction
The paper asks whether GNN improvements found on small, narrow molecular datasets transfer to larger, more diverse systems. It develops GemNet-OC on OC20 and tests model choices across datasets, finding that the choices and their effects can differ substantially.
1 Introduction
- The study asks whether improvements demonstrated on small, limited datasets generalize to large and diverse molecular systems.The motivating dataset dimensions are chemical diversity, system size, dataset size, and test-set difficulty.Source: §1 · p. 2
- GemNet-OC outperforms the previous state of the art on OC20 by 16% while training 10 times faster than previous models.The model incorporates optimizations for quadruplet interactions and dihedral angles in large systems, plus an interaction hierarchy for long-range interactions.Source: §1 · p. 2
- Model components can have substantially different effects across datasets, and all four dataset properties can contribute to those differences.The comparisons include MD17, COLL, OC20, and six OC20 subsets isolating dataset properties.Source: §1 · p. 2
- OC-2M is identified as a smaller subset whose model trends correlate well with OC20, supporting cheaper model development.The paper provides OC-2M baseline results for future work.Source: §1 · p. 2
2 Background and related work
Molecular simulation models use atomic structures to predict energies and forces, while GNNs represent atoms and their interactions as graphs. Prior work motivates examining how model trends change across datasets.
Learning atomic interactions
- Molecular simulation models take atomic numbers and positions as inputs and predict system energy and atomic forces.Forces can drive simulations or geometric optimization; conservative forces are energy gradients and support energy conservation and path independence.Source: §2 · p. 3
Graph neural networks
- GNNs represent atomic systems as graphs, typically connecting atom pairs within a distance cutoff, and update information through message passing.Modern molecular examples include SchNet, PhysNet, Cormorant, DimeNet, PaiNN, SpookyNet, and SpinConv.Source: §2 · p. 3
- GemNet extends two-level message-passing schemes to a multi-level interaction hierarchy using both edge and atom embeddings.The related approaches include MXMNet's two-level scheme and Alon and Yahav's global aggregation mechanism.Source: §2 · p. 3
Model trends between datasets
- Prior findings on scaling with training-set size motivate testing whether model choices on subsets transfer, but model trends can differ across datasets.Differences have also been observed for subsets of the same dataset beyond 200,000 samples.Source: §2 · p. 3
3 Datasets
The paper characterizes molecular datasets by their complexity, size, and test-set domain shift, then constructs OC20 subsets to isolate these factors. It argues that dataset properties affect how model choices transfer between benchmarks.
Datasets for molecular simulation
- OC20 comprises adsorbates binding to catalyst surfaces in periodic cells, and the paper compares it with other molecular datasets used for simulation.Table 1 summarizes dataset coverage, system sizes, sample counts, and test-set correlations.Source: §3 · pp. 3–4, Table 1 · p. 4
Dataset aspects
- The paper defines four dataset aspects: chemical diversity, system size, dataset size, and test-set domain shift.These properties affect data complexity, available learning information, and the difficulty of generalizing to test samples.Source: §3 · p. 4
Data complexity
- OC20 has 70× more neighbor pairs than other datasets, although element and neighbor-pair counts are imperfect proxies for chemical diversity.COLL has only three elements but samples more conformational space than QM7-X; SchNet's force MAE is 172 eV/Å on COLL and 53.7 eV/Å on QM7-X.Source: §3 · p. 4
Dataset size
- MD17 contains 3,611,115 samples, but recent work typically trains on 1,000 samples per molecule for its simple near-equilibrium task.The paper reports eight molecule trajectories on average and notes that this training subset is less than 1% of the dataset.Source: §3 · pp. 4–5
Domain shift
- Many datasets use correlated test samples from the same trajectories, while OC20 includes out-of-distribution test sets with unseen adsorbates, catalysts, or both.The paper notes that MD17 tests minor conformational-space generalization on the same molecule and trajectory, limiting the significance of its results.Source: §3 · p. 5
OC20 subsets
- Six OC20 subsets isolate dataset aspects: OC-Rb and OC-Sn restrict catalyst diversity, OC-sub30 restricts system size, and OC-200k and OC-2M reduce sample count.OC-XS combines a 30-independent-atom limit with rubidium catalysts; the subsets also use specified OOD validation splits, and OC-XS has a same-trajectory split.Source: §3 · p. 5
4 GemNet-OC
GemNet-OC builds on GemNet's atom- and directed-edge embeddings and geometric interactions. Its design changes target throughput and memory control while retaining interaction information at multiple ranges.
GemNet-OC
- GemNet-OC represents atoms and directed edges, updating both with geometric information from distances, angles, and dihedral angles.The model extends GemNet's message-passing architecture.Source: §4 · p. 5
Neighbors instead of cutoffs
- The nearest-neighbor graph gives the same or better accuracy, triples GemNet-OC throughput, and fixes neighborhood size for more predictable compute and memory use.It replaces a fixed distance cutoff, which can disconnect some graphs or make others excessively dense across OC20's chemical diversity.Source: §4 · p. 6
Simplified basis functions
- Simplified basis functions increase throughput by 29% without hurting accuracy.The design decouples radial and angular bases and uses an outer product of order-0 spherical harmonics, which simplify to Legendre polynomials.Source: §4 · pp. 6–7
Tractable quadruplet interactions
- Quadruplet interactions use kqint = 8 rather than the larger embedding cutoff because their benefits flatten as more neighbors are included.These optimizations reduce their overhead to 31%, instead of the original 330%.Source: §4 · p. 7
Interaction hierarchy
- The interaction hierarchy combines expensive short-distance quadruplet interactions with cheaper medium-distance edge-to-edge and longer-range atom-to-atom messages.Atom-to-atom messages use a 12 Å cutoff without neighbor restrictions; atom-to-edge and edge-to-atom interactions each add roughly 10% overhead.Source: §4 · p. 7
Further architectural improvements
- Further architectural changes combine outputs from interaction blocks, improve atom embeddings, and add atom embeddings directly to energy output blocks.These changes support information combination and flow with less than 2% computational overhead.Source: §4 · p. 7
5 Results on OC20
On OC20, GemNet-OC variants set new benchmarks across all three tasks, with GemNet-OC-L outperforming prior models by 16%. GemNet-OC also performs strongly with OC-2M training data and converges efficiently.
OC20 tasks
- OC20 evaluation covers S2EF, IS2RS, and IS2RE; this work focuses on relaxation-based IS2RE centered on S2EF data and models.Table 2 includes S2EF validation results and test results for all three tasks.Source: §5 · p. 7
Results
- 16%: GemNet-OC-L outperforms prior models on OC20, and GemNet-OC variants set the state of the art across all three tasks.Table 2 reports results averaged across all four OC20 splits.Source: Table 2 · p. 8, §5 · p. 8
- GemNet-OC trained on OC-2M outperforms pre-GemNet models trained on full OC20 and performs on par with GemNet-dT using 70 times less training data.OC20 contains approximately 134 million samples.Source: §5 · p. 8, Table 2 · p. 8
Training time
- After 600 GPU hours, GemNet-OC surpasses GemNet-dT; Figure 2 also reports that it is ∼12x faster than GemNet-XL.GemNet-OC-L reaches GemNet-XL’s final accuracy in 10 times fewer GPU hours and continues to improve.Source: Figure 2 · p. 8, §5 · p. 8
6 Model trends across datasets
Model choices often have different effects across datasets, reflecting differences in chemical diversity, domain shift, and training-set size. OC-2M most closely tracks OC20 among the tested subsets, while smaller or less diverse subsets can yield different trends.
Model width and depth
- Model width and depth show dataset-specific optima rather than steadily improving with size, with low-width optima in several smaller or less representative datasets.Increasing model size brings only modest accuracy gains, especially on out-of-distribution data; the result may be affected by model initialization.Source: §6 · p. 9
Batch size
- Most datasets favor batch sizes around 32 to 128, but MD17 and out-of-distribution validation sets favor particularly low optimal batch sizes.Convergence speed is fairly consistent across batch sizes when measured by samples seen rather than gradient steps.Source: §6 · p. 9
Neighbors and cutoffs
- Cutoff and neighbor trends are mostly consistent but vary in strength: fully connected graphs help MD17 and COLL, while nearest-neighbor graphs outperform cutoffs on OC20.Increasing quadruplets to 20 brings only a small improvement and about 20% runtime overhead.Source: §6 · pp. 9–10, Figure 4 · p. 10
Basis functions
- Bessel bases perform better on datasets with small chemical diversity but worse on OC20; multi-order Bessel bases are slower, while Legendre outer products match spherical harmonics faster.The comparison includes 6 and 32 Bessel functions against 128 Gaussian basis functions.Source: §6 · p. 10
Ablation studies
- Ablations show opposite effects across datasets; symmetric message passing and quadruplet interactions are among the most consistent choices.Additional interactions have little impact on in-distribution OC20 but benefit OC20 out-of-distribution data; atom-edge and edge-atom interactions also accelerate convergence.Source: §6 · pp. 10–11, Figure 5 · p. 10, Figure 6 · p. 11
Correlation between datasets
- OC-2M gives the most similar model-choice results to OC20; comparable difficulty and complexity depend on chemical diversity, domain shift, and sufficient training-set size.Uniformly sampled OC-2M preserves OC20’s chemical diversity and test difficulty; reducing the subset to OC-200k breaks down this similarity.Source: §6 · pp. 11–12
7 Conclusion
The paper studies how developing GNNs on large datasets differs from relying on small benchmarks. It finds that model choices vary with dataset properties, while OC-2M provides a correlated proxy for OC20 at lower computational cost.
Conclusion
- GemNet-OC is substantially faster to train than previous models and achieves state-of-the-art results on all OC20 tasks.The model was developed using the large OC20 dataset.Source: §7 · p. 12
- Model components and hyperparameters can have disparate or opposite effects across datasets, including within force prediction.The paper links consistent model choices to comparable chemical diversity, domain shift, and sufficient training-set size.Source: §7 · p. 12
- Uniformly sampling a sufficiently large subset produces OC-2M, a well-correlated proxy that supports model development for OC20 at a fraction of its computational cost.The authors provide baseline results for OC-2M and caution that model choices can vary when relevant dataset properties change.Source: §7 · p. 12
A Training and hyperparameters
The appendix describes shared training choices and dataset-specific adjustments. Models use a joint energy-and-force loss with direct force predictions for the ablations, while gradient-based force training is unstable on OC20.
Training setup
- All models use the same base hyperparameters, but effective batch sizes vary across dataset ablations to reach convergence in reasonable time.OC and non-OC datasets use different learning-rate schedules and evaluation intervals, and training hardware varies by dataset.Source: A Training and hyperparameters · p. 16
Target scaling
- OC energies are standardized by subtracting the dataset mean and dividing by its standard deviation, and forces use the same energy standard deviation.For MD17 and COLL, only the mean energy is subtracted.Source: A Training and hyperparameters · p. 16
Loss and force prediction
- The training objective combines energy and force errors with coefficients λ and ρ, and all model ablations make direct force predictions.The learnable energy and force functions share parameters; MD17 and COLL often perform better with gradient-based force predictions.Source: A Training and hyperparameters · p. 16
Optimization
- Gradient-based force models on OC20 often encounter numerical instability and NaNs early in training; models use AMSGrad and train to a maximum epoch limit or until learning-rate steps are exhausted.The stopping condition is whichever occurs first.Source: A Training and hyperparameters · p. 16
B Additional experimental results
The additional experiments show that early training can mislead model selection, and that trends across datasets and validation splits are not interchangeable. The appendix supplies complementary out-of-distribution comparisons and test-result tables.
Early stopping
- Early model performance can mislead: OC20 results become well-correlated with final accuracy only late in training, so early stopping is not a reliable shortcut for model development.The preferred basis can also reverse from early training to convergence; OC-2M behaves more consistently over training.Source: B Additional experimental results · p. 16, Figure 8 · p. 18, Figure 9 · p. 18
Basis projection sizes
- Larger basis projection sizes have a minor benefit on OC20 but are detrimental on almost all smaller datasets.Figure 10 reports their impact on force MAE.Source: Figure 10 · p. 18
OOD validation results
- On the OOD validation set, model-size optima appear on most datasets, while OC-XS, OC-Rb, and OC-Sn use in-distribution catalysts.The figure compares batch size, edge embedding size, and number of blocks.Source: Figure 11 · p. 19
- On OOD validation, the Bessel basis performs better on low-diversity datasets but worse on OC20, consistent with in-distribution trends.The figure compares cutoffs, neighbors, and basis functions.Source: Figure 12 · p. 19
- OOD ablations show that some changes, including symmetric message passing and quadruplet interaction, are fairly consistent across OC20 datasets, while others vary.The figure reports force MAE for proposed GemNet and GemNet-OC components.Source: Figure 13 · p. 19
- OC-2M is the most correlated with OC20 across model force MAEs and validation sets, including out-of-distribution comparisons.The appendix also provides OC20 test results for each test split and supplementary OOD validation plots.Source: B Additional experimental results · p. 16, Figure 14 · p. 20