ForceNet: A Graph Neural Network for Large-Scale Quantum Calculations

TL;DR

This work carefully design a scalable and expressive GNN model, ForceNet, and applies it to OC20 (Chanussot et al., 2020), an unprecedentedly-large dataset of quantum physics calculations, able to predict atomic forces more accurately than state-of-the-art physics-based GNNs while being faster both in training and inference.

Abstract

With massive amounts of atomic simulation data available, there is a huge opportunity to develop fast and accurate machine learning models to approximate expensive physics-based calculations. The key quantity to estimate is atomic forces, where the state-of-the-art Graph Neural Networks (GNNs) explicitly enforce basic physical constraints such as rotation-covariance. However, to strictly satisfy the physical constraints, existing models have to make tradeoffs between computational efficiency and model expressiveness. Here we explore an alternative approach. By not imposing explicit physical constraints, we can flexibly design expressive models while maintaining their computational efficiency. Physical constraints are implicitly imposed by training the models using physics-based data augmentation. To evaluate the approach, we carefully design a scalable and expressive GNN model, ForceNet, and apply it to OC20 (Chanussot et al., 2020), an unprecedentedly-large dataset of quantum physics calculations. Our proposed ForceNet is able to predict atomic forces more accurately than state-of-the-art physics-based GNNs while being faster both in training and inference. Overall, our promising and counter-intuitive results open up an exciting avenue for future research.

Figures (7)

• Figure 1: Comparison of S2F (atomic force prediction) performance across different models, while taking computational efficiency into account. (Left): Comparison of validation learning curves, where $x$-axis is training GPU days in the log-scale (lower left is better). (Right): Comparison of validation performance and inference time in GPU hours, measured over the in-distribution validation set (lower left is better). GPUs with the same specs are used for fair comparison (details in Section \ref{['subsec:model_config']}).
• Figure 2: Illustration of sampled systems from the OC20 dataset OC20. Each system consists of adsorbate (the small molecule on the surface) and catalysis (the large grid-like molecule sitting below the adsorbate), and is repeated in the direction of the horizontal axes infinitely. Our ForceNet aims to efficiently predict per-atom forces.
• Figure 3: Model diagram for messages $\bm{m}_{st}$ (from atom $s$ to atom $t$) used by ForceNet in Eqns. (\ref{['eq:model_message']}) and (\ref{['eq:conditional_filter']}). The key components are (a) the expressive conditional filter $\bm{F}_c$ that is dependent on full edge feature $\bm{e}_{st}$ (complete 3D relative placement information) as well as source and target node embeddings, $\bm{h}^{(k)}_s$ and $\bm{h}^{(k)}_t$, (b) the basis function $\bm{B}$ over the edge feature that helps the network to accurately capture atomic interactions, and (c) the smooth curved non-linearity of the Swish activation.
• Figure 3: Ablations on the architecture of ForceNet.
• Figure 4: Ablations on basis and activation functions in ForceNet.