Improving Molecular Modeling with Geometric GNNs: an Empirical Study

TL;DR

The paper conducts an empirical study of Geometric GNNs for 3D atomic systems, probing how canonicalization, graph creation, and auxiliary tasks affect predictive performance and symmetry handling on OC20 and QM9 tasks. It finds that approximate canonicalization methods like SFA can rival or exceed exact $E(3)$-equivariant approaches in practice, and that graph construction is robust across a range of cutoffs, with graph rewiring offering substantial scalability gains. Noisy Nodes emerge as a powerful auxiliary task enabling much deeper networks to outperform shallower baselines, while pre-training on larger, related tasks shows transferable benefits though final convergence closely tracks direct training. Overall, the study provides practical guidance for selecting modeling components and highlights promising directions in transfer learning and multi-task architectures for molecular property prediction.

Abstract

Rapid advancements in machine learning (ML) are transforming materials science by significantly speeding up material property calculations. However, the proliferation of ML approaches has made it challenging for scientists to keep up with the most promising techniques. This paper presents an empirical study on Geometric Graph Neural Networks for 3D atomic systems, focusing on the impact of different (1) canonicalization methods, (2) graph creation strategies, and (3) auxiliary tasks, on performance, scalability and symmetry enforcement. Our findings and insights aim to guide researchers in selecting optimal modeling components for molecular modeling tasks.

Figures (7)

• Figure 1: Similarity matrix of the embeddings of the atoms of a randomly picked system for different interaction blocks from the training set of OC20. The same system is used every time to be able to compare the different results.
• Figure 2: MAD values of the graph embeddings (averaged over 50 randomly sampled graphs of the train set) throughout the interaction layers for various models. "FAENet top" models are trained with the top configs of the classical FAENet model duval2023faenet but with more epochs and lower batch size (128). "FAENet aux" models are our models trained on IS2RE with IS2RS auxiliary task. A model having xx interaction layers is indicated as "XXi".
• Figure 3: Ewald Summation interactions. The interaction term (left) is the result of the short-range interaction (middle) and the long-term interaction (right) which are both computed using cutoffs on respectively the real and the Fourier space. Adapted from kosmala2023ewald
• Figure 4: Similarity matrix of the embeddings of the atoms of a system for different interaction blocks on FAENet and SchNet with Ewald message passing on the models. The visualized layers here are the standard interaction blocks in the first row and the Ewald interaction blocks in the second row. The two outputs are summed to get the final representation for Ewald shown in Figure \ref{['fig:ewald_embeddings']}.
• Figure 5: Same plots as Figure \ref{['fig:ewald_embeddings']} but with a second randomly picked system from the OC20 train split.