SlowGuess/NeurIPS_2023_OCR huggingface dataset:
QuinNet: Efficiently Incorporating Quintuple Interactions into Geometric Deep Learning Force Fields
Zun Wang
Microsoft Research AI4Science
Beijing, China, 100084
zunwang@microsoft.com
Yichi Zhou
Guoqing Liu
Tong Wang
Bin Shao
binshao@microsoft.com
Machine learning force fields (MLFFs) have instigated a groundbreaking shift in molecular dynamics (MD) simulations across a wide range of fields, such as physics, chemistry, biology, and materials science. Incorporating higher order many-body interactions can enhance the expressiveness and accuracy of models. Recent models have achieved this by explicitly including up to four-body interactions. However, five-body interactions, which have relevance in various fields, are still challenging to incorporate efficiently into MLFFs. In this work, we propose the quintuple network (QuinNet), an end-to-end graph neural network that efficiently expresses many-body interactions up to five-body interactions with ab initio accuracy. By analyzing the topology of diverse many-body interactions, we design the model architecture to efficiently and explicitly represent these interactions. We evaluate QuinNet on public datasets of small molecules, such as MD17 and its revised version, and show that it is compatible with other state-of-the-art models on these benchmarks. Moreover, QuinNet surpasses many leading models on larger and more complex molecular systems, such as MD22 and Chignolin, without increasing the computational complexity. We also use QuinNet as a force field for molecular dynamics (MD) simulations to demonstrate its accuracy and stability, and conduct an ablation study to elucidate the significance of five-body interactions. We open source our implementation at https://github.com/Zun-Wang/QuinNet.
Machine learning force fields (MLFFs) [1, 2, 3, 4, 5, 6] have brought about a transformative shift in molecular dynamics (MD) simulations across diverse fields, including physics, chemistry, biology, and material science. MLFFs enable superior accuracy and reliability by leveraging ab initio and experimental data to learn interatomic interactions, while simultaneously keeping the low cost of empirical force fields. As a result, MLFFs provide a powerful tool for accurately modeling molecular
interactions, and have the potential to revolutionize our understanding of complex systems at both the atomistic and macroscopic scales.
Recent geometric deep learning potentials, represent atoms in a system as nodes and chemical bonds as edges in a graph while incorporating physical symmetries. These methods have been proven to be data-efficient and effective in modeling many-body interactions [7]. The development track of geometric deep learning potentials is similar to that of empirical force fields, which follow the many-body expansion theory [8, 9, 10]. Previous studies [11, 12, 13, 14, 15, 16, 17, 18, 19, 20] have primarily adopted a hierarchical approach to model truncated energy correction terms, enabling the incorporation of higher-order interactions. As higher-order many-body interactions play a critical role in accurately modeling complex molecular systems, this approach has shown promise in leveraging graph neural networks (GNNs) [21] to construct more expressive and accurate models.
Nowadays, GNN models have incorporated many-body interactions up to four-body interactions explicitly [17, 18, 19, 20] and incorporating higher order many-body interactions is still a challenge. The neglect of higher-order many-body interactions may result in incomplete molecular representations [22] and limited accuracy in certain cases [23, 24]. For example, five-body interactions play an important role in various fields such as coarse-grained protein force field [25], organic molecules [26], crystal vibrations [27, 28], electrostatic interaction potentials [29], and so on.
In this paper, we develop a MLFF based on GNNs, called quintuple network (QuinNet), to address the challenge of incorporating higher order many-body interactions, up to quintuple, into molecular representations without introducing more computational complexity. To achieve this, we analyze the topology of diverse many-body interactions and design QuinNet to explicitly represent these interactions efficiently. Our approach outperforms existing methods in terms of accuracy on several large molecular systems (e.g., MD22 and Chignolin dataset) while remaining comparable on small molecules (e.g., MD17 and revised MD17 dataset). Additionally, computational complexity analysis demonstrates the efficiency of our model. Our work has the potential to advance the field of MLFFs by providing methods to incorporate higher order many-body interactions.
The field of MLFFs has witnessed significant progress since the introduction of the first neural network-based framework by Behler and Parrinello [30]. However, the inherent indistinguishability of atoms has posed a challenge in ordering them, which has limited the success of this approach. To overcome this challenge, GNNs have emerged as a promising solution, leveraging the power of graph formalism to focus on relationships among entities rather than individual node properties.
Currently, two mainstream GNN-based methods have been developed for constructing force fields: group theory-based methods and direction-based methods. Group theory-based methods preserve physical symmetries by imposing trainable weights as the representation of related groups [22, 31]. These equivariant models, including Tensor Field Networks [32], Cormorant [33], SE(3)-Transformers [34], NequIP [7], Equiformer [35], and Allegro [36], have demonstrated high performance across diverse tasks. Furthermore, the MACE [37], which was developed hierarchically using a body order expansion manner, has been proposed as a unifying framework of E(3)-equivariant atom-centered interatomic potentials, extending the ACE [38] framework to include methods built on equivariant message passing neural networks (MPNNs).
On the other hand, direction-based methods, such as SchNet [11, 12], DimeNet (DimeNet++) [13, 14], PaiNN [15], ET [16], GemNet [17], SphereNet [18], ComENet [19], and ViSNet [20], explicitly model interatomic interactions by incorporating more geometric information. This progress has motivated us to investigate the potential performance enhancements of explicitly incorporating higher-order interactions into GNN-based force fields.
3.1 GNN-based force fields
The key idea of graph neural networks boils down to a message-passing framework [39], wherein each node iteratively updates its embedding by aggregating messages from its neighboring nodes.
With  $h_v$  and  $e_{vw}$  denoting node features and edge features respectively, the message-passing layer is defined as
$$
m _ {v} ^ {t + 1} = \sum_ {w \in \mathcal {N} (v)} M _ {t} \left(h _ {v} ^ {t}, h _ {w} ^ {t}, e _ {v w} ^ {t}\right), \tag {1}
$$
$$
h _ {v} ^ {t + 1} = U _ {t} \left(h _ {v} ^ {t}, m _ {v} ^ {t + 1}\right),
$$
where  $M_{t}$  and  $U_{t}$  are a message function and an update function, respectively. In the realm of GNN-based force field, physical symmetries such as translation, rotation, and permutation symmetries have been demonstrated to be intrinsic and highly effective for model design.
Group equivariance Let  $\mathcal{L}:\mathcal{X}\to \mathcal{Y}$  be a function that maps from input space  $\mathcal{X}$  to output space  $\mathcal{V}$ . We say that  $\mathcal{L}$  is equivariant with respect to a group  $G$  if for all group elements  $g\in G$
$$
\mathcal {L} \circ D ^ {\mathcal {X}} (g) = D ^ {\mathcal {Y}} (g) \circ \mathcal {L}, \tag {2}
$$
where  $D^{\mathcal{X}}$  is the representation of group  $G$  in the input vector space  $\mathcal{X}$ .
Group equivariant convolutional layer In recent years, there has been a growing interest in designing group equivariant neural networks due to the importance of respecting relevant symmetries as an inductive bias. Specifically, the group equivariant convolutional layer is defined as follows:
$$
\mathcal {L} _ {a c m _ {o}} ^ {(l _ {o})} \left(\vec {r} _ {a}, V _ {a c m _ {i}} ^ {l _ {i}}\right) := \sum_ {m _ {f}, m _ {i}} C _ {\left(l _ {f}, m _ {f}\right) \left(l _ {i}, m _ {i}\right)} ^ {\left(l _ {o}, m _ {o}\right)} \sum_ {b \in S} F _ {c m _ {f}} ^ {\left(l _ {f}, l _ {i}\right)} \left(\vec {r} _ {a b}\right) V _ {b c m _ {i}} ^ {l _ {i}}, \tag {3}
$$
where the node features are denoted by  $V$ , and  $\vec{r}_{ab}$  represents the relative position vector between point  $a$  and point  $b$ . The subscripts  $c, i, f,$  and  $o$  denote channel, input, filter, and output, respectively. Additionally,  $l$  and  $m$  correspond to the orbital quantum number and magnetic quantum number, respectively. This definition of the group equivariant convolutional layer is an important contribution to the development of neural networks that can effectively handle physical symmetries.
3.2 Empirical force fields
In classical molecular dynamics simulations, force fields play a crucial role in describing the potential energy  $\mathcal{V}(\mathbf{r}^N)$  of a system, where  $\mathbf{r}^N$  represents the positions of  $N$  particles. A common functional form for force fields is  $\mathcal{V}(\mathbf{r}^N) = E_{\mathrm{bonded}} + E_{\mathrm{non -bonded}}$ . The bonded energy term captures the deviation of bonds and angles from their equilibrium values, while the non-bonded interactions include long-range forces, such as electrostatic and van der Waals interactions. To represent the bonded interactions hierarchically, bond stretching, angle bending, bond rotation, and other factors can be taken into account.
Improvement of torsion potential The torsion potential is a crucial aspect of molecular modeling, representing the energy associated with the rotation of dihedral angles between bonded atoms. While many of the torsion terms are comprised of only one term from the cosine series expansion [40], it has been discovered that including more than one term is necessary for certain bonds, such as gauche conformations [41, 42]. This involves the use of higher order cosine series, represented by  $\sum_{n=0}^{N} \frac{V_n}{2} [1 + \cos(n\omega - \gamma)]$ , where  $\omega$  denotes the torsion angle and  $V_n$  is the barrier height.
Improper torsions and out-of-plane bending motions Experimental observations have shown that in certain molecular systems, the atom connected to a ring tends to remain in the same plane due to the  $\pi$  bonding energy. However, the force field discussed in the previous section predicts an equilibrium structure with the atom out of plane. In order to correct this discrepancy, an addi
Figure 1: Two methods to represent out-of-plane bending terms: (a) calculating the angle between a bond from the central atom and the plane defined by the central atom and the other two atoms and (b) calculating the height of the central atom above a plane defined by the other three atoms.
tional term or terms need to be incorporated into the force field to ensure that the  $\mathrm{sp}^2$  carbon and its three bonded atoms remain in the same plane. One simple approach is to include an "out-of-plane"
bending term [43]. Fig. 1 illustrates two methods that can be used to model the out-of-plane bending contributions. The first method involves calculating the angle between a bond from the central atom and the plane defined by the central atom and the other two atoms, as shown in Fig. 1 (a). The second method calculates the height of the central atom above a plane defined by the other three atoms, as shown in Fig. 1 (b).
3.3 Challenge of incorporating higher order many-body interactions into GNNs explicitly
Higher-order many-body interactions could be learned implicitly by stacking more MPNN layers, while incorporating these interactions explicitly into GNNs would improve the performance significantly. Incorporating many-body interactions explicitly into GNN models heavily relies on the selection of appropriate physical quantities, such as bond lengths, bond angles, and dihedral angles. However, for higher order many-body interactions, finding a suitable quantity poses a significant challenge. Moreover, calculating such a physical quantity increases the computational complexity as the order of interactions increases. Therefore, solving this problem while simultaneously reducing computational complexity is of utmost importance.
To explicitly incorporate five-body interactions into GNNs, we first analyze the topology of three-, four-, and five-body interactions to identify the key physical quantities required to describe these many-body interactions. Based on this analysis, we design the model architecture to effectively capture and represent these interactions.
4.1 Topology of five-body interactions
Five-body interactions have been shown to play a critical role in various cases, such as reproducing specific phenomena in protein systems [25]. While three-body and four-body interactions can be effectively captured by incorporating angular and dihedral angular information, which physical quantity is suitable for representing the five-body interactions explicitly is still a question.
In this work, we propose to use dihedral angles as a unified physical quantity to describe five-body interactions. Fig. 2 illustrates the topology of three-body, four-body, and five-body interactions, respectively. Fig. 2 (a) and (c) show that three-body interactions and improper torsion in four-body interactions could be represented as angles; and Fig. 2 (b) implies that torsion in four-body interactions could be represented as dihedral angles. There are only three different topologies for five-body interactions (Fig. 2 (d)-(f)) and all of them could be represented as dihedral angles. To calculate dihedral angles, we first compute the cosine of the angle between two planes formed by vectors connecting neighboring atoms. For instance, in Fig. 2 (d), atoms  $j_{1}, j_{2}, j_{3}$ , and  $j_{4}$  form a neighborhood around atom  $i$ , and the normal vector of the plane composed of vectors  $\vec{r}_{ij_1}$  and  $\vec{r}_{ij_2}$  can be computed using their outer product, i.e.  $(\vec{r}_{ij_1} \times \vec{r}_{ij_2}) \cdot (\vec{r}_{ij_3} \times \vec{r}_{ij_4})$ . The dihedral angle
Figure 2: A schematic diagram that describes the topology of (a) three-body interaction (angles), (b) four-body interaction (torsions), (c) four-body interaction (improper torsions), and (d-f) five-body interactions. The marker  $\odot$  on a plane represents the normal vector of this plane.
between these two planes can be easily determined as the supplementary angle to the angle between their normal vectors, which can be calculated by taking their inner product. Further details on the relationships between different many-body interactions refer to the Supplementary Materials.
4.2 Architecture of QuinNet
Inspired by the topology analysis of many-body interactions, we design the model architecture of QuinNet. The overall schematic diagram of QuinNet is shown in Fig. 3 (a), which explicitly incorporates many-body interactions and reduces computational complexity simultaneously. In particular, we design the QuinNet module, as shown in Fig. 3 (b), where different many-body (i.e., from 3-body to 5-body) interactions are handled separately and linearly combined as the output of the module.
Three-body interactions The three-body interaction, which is related to angular information as depicted in Fig. 2 (a), is handled using the PaiNN [15] approach (refer to Fig. 3 (c)). The equation for the three-body interaction is expressed as:
$$
\left\| \sum_ {j \in \mathcal {N} _ {i}} \hat {r} _ {i j} \right\| ^ {2} = \sum_ {j, k \in \mathcal {N} _ {i}} \langle \hat {r} _ {i j}, \hat {r} _ {i j} \rangle = \sum_ {j, k \in \mathcal {N} _ {i}} \cos \alpha_ {j i k}, \tag {4}
$$
where  $\hat{r}_{ij}$  denotes the normalized vector of relative position between atom  $i$  and  $j$ ,  $\vec{r}_{ij}$ , and  $\langle \cdot, \cdot \rangle$  is inner product. This approach incorporates the representation of bond angular information into node  $i$ .
Four-body interactions (torsion) The incorporation of torsion interactions in force fields necessitates the use of dihedral angular terms, as depicted in Fig. 2 (b). To address this, we adopt the ViSNet [20] model, which explicitly accounts for torsion energy, as shown in Fig. 3 (b). This is achieved through vector rejection of the direction unit of node  $i$ ,  $\vec{v}_i$ , which is equivalent to Gram-Schmidt normalization. In 3-dimensional space, the process reduces to calculating the cross product of the normal vector of two planes, as what we summarized in the last section. The torsion angles are then represented via
$$
\left(\sum_ {j \in \mathcal {N} _ {i}} \hat {r} _ {i j} \times \hat {r} _ {i j _ {1}}\right) \cdot \left(\sum_ {k \in \mathcal {N} _ {j _ {1}}} \hat {r} _ {j k} \times (- \hat {r} _ {i j _ {1}})\right) = \sum_ {\substack {j \in \mathcal {N} _ {i}, \\ k \in \mathcal {N} _ {j _ {1}}}} \langle \vec {n} _ {i j _ {1} j _ {2}}, \vec {n} _ {i j _ {1} k} \rangle , \tag{5}
$$
where  $\vec{n}_{ijk_1}$  is the normal vector of the plane composed of nodes  $i, j$ , and  $k_{1}$  ( $i, j$ , and  $k_{1}$  are non-colinear). This approach constructs hidden states with dihedral angular information for edge  $e_{ij}$ .
Four-body interactions (improper torsion) As ViSNet does not incorporate improper torsion term, the QuinNet module incorporates the out-of-plane bending term by computing it in the same manner as Fig. 2 (c), as depicted in Fig. 3 (e). Specifically, the term is computed using the equation
$$
\sum_ {j \in \mathcal {N} _ {i}} \hat {r} _ {i j} \cdot \left[ \left(\sum_ {j \in \mathcal {N} _ {i}} \alpha_ {j} \hat {r} _ {i j}\right) \times \left(\sum_ {j \in \mathcal {N} _ {i}} \beta_ {j} \hat {r} _ {i j}\right) \right] = \sum_ {j _ {1}, j _ {2}, j _ {3} \in \mathcal {N} _ {i}} \gamma_ {i j _ {1} j _ {2}} \langle \vec {r} _ {i j _ {3}}, \vec {n} _ {i j _ {1} j _ {2}} \rangle . \tag {6}
$$
where the different coefficients  $\alpha_{j}$  and  $\beta_{j}$  are introduced to prevent the term from being zero due to the properties of the cross product and the  $\gamma$  is the coefficient from the product of  $\alpha$  and  $\beta$ . The out-of-plane bending term is an essential component of nodes' embedding.
Five-body interactions@I The evaluation the five-body@I term could be represented as dihedral angles (depicted in Fig. 2 (d) and Fig. 3 (f)):
$$
\left\| \left(\sum_ {j \in \mathcal {N} _ {i}} \alpha_ {j} \hat {r} _ {i j}\right) \times \left(\sum_ {j \in \mathcal {N} _ {i}} \beta_ {j} \hat {r} _ {i j}\right) \right\| ^ {2} = \sum_ {j _ {1}, j _ {2}, j _ {3}, j _ {4} \in \mathcal {N} _ {i}} \gamma_ {i j _ {1} j _ {2} j _ {3} j _ {4}} \langle \vec {n} _ {i j _ {1} j _ {2}}, \vec {n} _ {i j _ {3} j _ {4}} \rangle . \tag {7}
$$
As noted before, different coefficients  $\alpha_{j}$  and  $\beta_{j}$  are assigned to avoid the term becoming zero and the term provides a node embedding. Moreover, the term encompasses a portion of the description for improper interactions, and further details can be found in the Supplementary Materials.
Figure 3: A schematic diagram of QuinNet. (a) The overall architecture of QuinNet, where  $\{Z\}$  represents the set of atomic numbers. The initial node features are zero vectors of a fixed dimension. The QuinNet module inputs scalar and vectorial node features  $s_i$  and  $\vec{v}_i$ , radial basis function (RBF) as scalar edge feature  $e_{ij}$ , and relative position vector  $\vec{r}_{ij}$  as vectorial edge feature, and then outputs an updated scalar and vectorial node feature and scalar edge feature for the next layer. The set of spherical harmonic functions,  $Y_{l}m$ , is denoted as  $Y_{l}$ . The final scalar node features are summed to form the graph feature, and the target prediction is obtained. (b) The QuinNet module consists of six modules that handle (c) 3-body interactions, (d) torsion terms in 4-body interactions, (e) improper terms in 4-body interactions, and (f)-(h) five-body interactions, respectively. The normal vector corresponding to node  $j$  is represented by  $\vec{n}_j$ . In a linear layer, the weights are referred to as  $W$ , while the term  $b$  signifies the bias.
Five-body interactions@II To compute the five-body interaction@II term (refer to Fig. 2 (e) and Fig. 3 (g)), we use the following equation to construct hidden states on edges:
$$
\left. \left\| \left(\sum_ {k \in \mathcal {N} _ {j _ {1}}} \alpha_ {k} \hat {r} _ {k j _ {1}} \times \sum_ {k \in \mathcal {N} _ {j _ {1}}} \beta_ {k} \hat {r} _ {k j _ {1}}\right) \right| _ {j _ {1} \in \mathcal {N} _ {i}} \right\| ^ {2} = \sum_ {\substack {k _ {1} \in \mathcal {N} _ {j _ {1}}, \\ k _ {2} \in \mathcal {N} _ {j _ {2}}}} \gamma_ {j _ {1} j _ {2} k _ {1} k _ {2}} \langle \vec {n} _ {i j _ {1} k _ {1}}, \vec {n} _ {i j _ {2} k _ {2}} \rangle \Bigg | _ {j _ {1}, j _ {2} \in \mathcal {N} _ {i}}, \tag{8}
$$
where  $\big|_{j} \in \mathcal{N}_{i}$  indicates that the values should be chosen according to the source nodes of node  $i$ .
Five-body interactions@III The five-body interaction@III term, as illustrated in Fig. 2 (f) and Fig. 3 (h), is computed as follows,
$$
\left. \left(\sum_ {j \in \mathcal {N} _ {i}} \alpha_ {j} \hat {r} _ {i j} \times \sum_ {j \in \mathcal {N} _ {i}} \beta_ {j} \hat {r} _ {i j}\right) \cdot \left(\sum_ {k \in \mathcal {N} _ {j}} \alpha_ {k} \hat {r} _ {j k} \times \sum_ {k \in \mathcal {N} _ {j}} \beta_ {k} \hat {r} _ {j k}\right) \right| _ {j \in \mathcal {N} _ {i}} = \sum_ {\substack {j 1, j 2 \in \mathcal {N} _ {i}, \\ k _ {1}, k _ {2} \in \mathcal {N} _ {j}}} \gamma_ {i j _ {1} j _ {2} k _ {1} k _ {2}} \langle \vec {n} _ {i j _ {1} j _ {2}}, \vec {n} _ {j k _ {1} k _ {2}} \rangle \Bigg | _ {j \in \mathcal {N} _ {i}},
$$
which constructs the representation of edges.
Higher order cosine series Since it is beneficial to include higher order cosine series terms in the force field [41, 42], we utilize the vector addition theorem of spherical harmonic functions. Although all the angular calculations in Eq. 4 - Eq. 9 only consider one order of cosine terms, the higher-order cosine terms could be incorporated easily by using the following equation:
$$
P _ {l} (\cos \theta) = \frac {4 \pi}{2 l + 1} \sum_ {m = - l} ^ {l} Y _ {l m} ^ {*} (\hat {u}) Y _ {l m} (\hat {v}), \tag {10}
$$
where  $\theta$  denotes the angle between two unit vectors  $\hat{u}$  and  $\hat{v}$ , and  $P_{l}$  and  $Y_{l}$  represent Legendre polynomials and spherical harmonic functions of  $l$ -th order, respectively.  $Y_{l}^{*}$  is the complex conjugate of  $Y_{l}$ . As  $l$  increases, higher-order cosine terms are incorporated into the calculation, enabling us to improve the accuracy of our model and make it suitable for a broader range of applications. Additional proof can be found in the Supplementary Materials.
4.3 Complexity analysis
Inspired by the design of PaiNN [15], QuinNet incorporates many-body interactions by calculating relevant physical quantities during the message passing layer, eliminating the need for calculations beforehand. The operations within the QuinNet module can be decomposed into cross product and inner product, which are independent of the number of nodes' neighbors. As a result, the computational complexity of each many-body interaction is significantly reduced to  $\mathcal{O}(|\mathcal{N}|)$ , where  $|\mathcal{N}|$  represents the number of neighbors for each node. This approach allows QuinNet to efficiently model complex interactions, enhancing its overall performance. A more comprehensive analysis refers to the Supplementary Materials.