Equivariant Graph Network Approximations of High-Degree Polynomials for Force Field Prediction

TL;DR

This work analyzes the equivariant polynomial functions for the equivariant architecture, and introduces a novel equivariant network, named PACE, which demonstrates state-of-the-art performance in predicting atomic energy and force fields, with robust generalization capability across various geometric distributions under molecular dynamics across different temperature conditions.

Abstract

Recent advancements in equivariant deep models have shown promise in accurately predicting atomic potentials and force fields in molecular dynamics simulations. Using spherical harmonics (SH) and tensor products (TP), these equivariant networks gain enhanced physical understanding, like symmetries and many-body interactions. Beyond encoding physical insights, SH and TP are also crucial to represent equivariant polynomial functions. In this work, we analyze the equivariant polynomial functions for the equivariant architecture, and introduce a novel equivariant network, named PACE. The proposed PACE utilizes edge booster and the Atomic Cluster Expansion (ACE) technique to approximate a greater number of $SE(3) \times S_n$ equivariant polynomial functions with enhanced degrees. As experimented in commonly used benchmarks, PACE demonstrates state-of-the-art performance in predicting atomic energy and force fields, with robust generalization capability across various geometric distributions under molecular dynamics (MD) across different temperature conditions. Our code is publicly available as part of the AIRS library https://github.com/divelab/AIRS/.

Figures (4)

• Figure 1: An architecture overview of PACE. A: The first message passing layer. The initial node features $\mathbf{x}_i^0$ and $\mathbf{x}_j^0$ are concatenated and transformed by an MLP. Then, a tensor product is applied to the transformed node features and edge spherical harmonics $Sph_{ij}$ with an RBF and MLP transformed edge distance as learnable weights. Next, its output is further combined with $Sph_{ij}$ using another tensor product, and the weights involve the initial node features and edge distance. These operations in the yellow box comprise the edge booster, aiming to enhance the model's expressiveness. The obtained edge-boosted messages are aggregated to form the atomic base. Then, an update module comprising a polynomial many-body interaction module and a skip connection is used to update the features of central nodes. B: The second message passing layer. A tensor product is applied to the edge message and updated node features that are obtained from the first layer. The learnable weights of the tensor product are based on the initial node features and distance of each neighboring pair. Then, edge messages are aggregated and node features are updated similarly to those in the first layer. C: An example of 4-body interaction in ACE. We aim to fit $\sum{\phi_{v_1}(\mathbf{r}_{ij_1})} \otimes \sum{\phi_{v_2}(\mathbf{r}_{ij_2})} \otimes \sum{\phi_{v_3}(\mathbf{r}_{ij_3})}$ using $A_{i v_1} \otimes A_{i v_2} \otimes A_{i v_3}$, where $\{\phi\}$ denotes the atomic base in ACE. D: Polynomial many-body interaction module. The atomic base $A_i$ is fed to multiple self-interaction layers separately to produce different $A_{iv}$. Then, tensor contraction is performed to produce $\tilde{a}_i$. E: Output. The invariant part of node features produced by both layers are transformed and summed to predict the deviation from the total molecular energy to its average.
• Figure 2: Illustration of tensor contraction in the polynomial many-body interaction module. In this figure, it demonstrates an example of 4-body interactions with $v_{\max}=3$ and final $L_3 = 0, M_3 = 0$. Note that the contract weights operation learns weighted summation over all paths $\eta[v]$, where $\eta[v] = (\ell_1, \ell_2, L_2, \cdots, \ell_v, L_v)$.
• Figure 3: Illustration of radial distribution functions (RDF) of MD trajectories. Values are averaged over five MD simulations with five initial molecular structures. The shell thickness $dr=0.05$ is used.
• Figure 4: Exploration of hyperparameter space.

Theorems & Definitions (8)

• Theorem 4.1
• Definition 4.2
• Definition A.1
• proof
• Lemma B.1
• proof
• Lemma B.2
• proof