Lightweight equivariant model for efficient machine learning interatomic potentials

TL;DR

A lightweight equivariant interaction graph neural network (LEIGNN) that can enable accurate and efficient interatomic potential and force predictions for molecules and crystals and achieves significant efficiency across diverse datasets.

Abstract

In modern computational materials science, deep learning has shown the capability to predict interatomic potentials, thereby supporting and accelerating conventional simulations. However, existing models typically sacrifice either accuracy or efficiency. Moreover, lightweight models are highly demanded for offering simulating systems on a considerably larger scale at reduced computational costs. Here, we introduce a lightweight equivariant interaction graph neural network (LEIGNN) that can enable accurate and efficient interatomic potential and force predictions for molecules and crystals. Rather than relying on higher-order representations, LEIGNN employs a scalar-vector dual representation to encode equivariant features. By learning geometric symmetry information, our model remains lightweight while ensuring prediction accuracy and robustness through the equivariance. Our results show that LEIGNN consistently outperforms the prediction performance of the representative baselines and achieves significant efficiency across diverse datasets, which include catalysts, molecules, and organic isomers. Furthermore, we conduct molecular dynamics (MD) simulations using the LEIGNN force field across solid, liquid, and gas systems. It is found that LEIGNN can achieve the accuracy of \textit{ab initio} MD across all examined systems.

Figures (5)

• Figure 1: Overview of GNNs applied to materials property prediction. (a) A typical GNN workflow is illustrated, emphasizing how GNNs can include a wider range of interactions (or message passing) by stacking multiple layers to extend the accessible radius. (b), (c), and (d) underscore the importance of integrating structural features into GNNs, which include: (b) distance only, (c) both distance and angles, and (d) all distance, angles, and dihedral angles. (e), (f), and (g) elucidate the concepts of invariance and equivariance within the context of energy and forces prediction.
• Figure 2: The overall architecture of LEIGNN. (a) LEIGNN uses message distributing, message passing, message updating, and message aggregating to iteratively update node representations. (b) LEIGNN consists of $T$ layers. (c) The message passing phase. (d) The message updating phase. The number of features after each operation is annotated in grey.
• Figure 3: An illustration of how vector representation evolves during message passing phase and message updating phase. (a) An example explains how to calculate the intermediate vector representation $\vec{\mathbf{m}}_1$. For the sake of simplicity, we ignore the first term of Eqn. \ref{['eqn:vec_update']}. The $f$-th vector in $\vec{\mathbf{m}}_{1}$ (i.e., $\vec{\mathbf{m}}_1^f$) can be interpreted as the force exerted by neighboring atom $j$ on atom $1$, where $w_{j1}^f$ is force magnitude and $\frac{\vec{\mathbf{r}}_{j1}}{\Vert \vec{\mathbf{r}}_{j1} \Vert}$ is the force direction. Subsequently, the total force exerted on atom $1$ is a linear combination of forces exerted on it by all other neighboring atoms $j$. Note that the forces are calculated $F$ times in parallel. (b) The update to $\vec{\mathbf{x}}_1^f$ is achieved through a linear combination of $F$ vectors within $\vec{\mathbf{m}}_1$. This computation is performed in parallel $F$ times to obtain $\vec{\mathbf{x}}_1$.
• Figure 4: Computational costs of the representative models (the lower the better for all metrics). (a) Training time. (b) Mean inference time. (c) Number of parameters. (d) FLOPs.
• Figure 5: MD simulations for LiPS, $\rm H_2O$, and $\rm CH_4$. (a) An overview of the three benchmark systems. (b) $h(r)$ for trajectories predicted by the LEIGNN. (c) RDF for trajectories predicted by LEIGNN.