Tensor improve equivariant graph neural network for molecular dynamics prediction

Abstract

Molecular dynamics(MD) simulations are essential for molecular structure optimization, drug-drug interactions, and other fields of drug discovery by simulating the motion of microscopic particles to calculate their macroscopic properties (e.g., energy). The main problems of the existing work are as follows: (1) Failure to fully consider the chemical bonding constraints between atoms, (2) Group equivariance can help achieve robust and accurate predictions of MD under arbitrary reference transformations and should be incorporated into the model design, (3) Tensor information such as relative position, velocity, and torsion angle can be used to enhance the prediction of molecular dynamics. And the existing methods are mainly limited to the scalar domain.In this paper, we propose a new model—tensor improve equivariant graph neural network for molecular dynamics prediction (TEGNN): (1) The model materialization of chemical bond constraints between atoms into geometric constraints. The molecule’s forward kinematic information (position and velocity) is represented by generalized coordinates. In this way, the interatomic chemical bonding constraints are implicitly and naturally encoded in the forward kinematics, (2) The equivariant information transfer is allowed in TEGNN, which significantly improves the accuracy and computational efficiency of the final prediction, (3) TEGNN introduces equivariant locally complete frames into scalar-only equivariant graph neural networks, thus allowing the projection of tensor information of a given order onto the frame. On the N-body dataset of simulated molecular systems consisting of particles, sticks, and hinges, as well as on the real dataset MD17 for molecular dynamics prediction, multiple experiments support that TEGNN is ahead of the current state-of-the-art GNN in terms of prediction accuracy, constraint satisfaction, and data efficiency.We extend the current state-of-the-art equivariant neural network model. The proposed TEGNN accommodates more tensor information and considers the chemical bonding constraints between atoms during motion, ultimately improving the performance of predicting the kinematic states of molecules.