Structural Causal Model for Molecular Dynamics Simulation

Qi Liu, Yuanqi Du, Fan Feng, Qiwei Ye, Jie Fu

Keywords: Molecular dynamics, causal discovery, neural relation inference

Abstract: Molecular dynamics (MD) simulations describe the mechanical behaviors of molecular systems through empirical approximations of interatomic potentials. Machine learning-based approaches can improve such potentials with better transferability and generalization. Among them, graph neural networks have prevailed as they incorporate the graph structure prior while learning the interatomic interactions. Nevertheless, the simple design choices and heuristics in devising graph neural networks make them lack an explicitly interpretable component to identify the true physical interactions within the underlying system. On the other extreme, physical models can give a rather comprehensive description of a system but are hard to specify. Causal modeling lies in between these two extremes, and can provide us with more modeling flexibility. In this paper, we propose a structural causal molecular dynamics model (SCMD), the first causality-based framework to model interatomic and dynamical interactions in molecular systems by inferring causal relationships among atoms from observational data. Specifically, we leverage the structural causal model (SCM) to model the interaction system of MD. To infer the SCM, we construct the graph in SCM as the dynamic Bayesian network (DBN), which is learned by a sequential generative model named SC-VAE. In the SC-VAE, the encoder and decoder infer the causal structure and temporal dynamics. All components are learned in an end-to-end fashion, and the DBN is learned in an unsupervised way. Furthermore, by concerning the underlying data generation process, inducing the causal structure and temporal dynamics of the system, one can enjoy a robust and flexible MD simulation model to explicitly capture the long-range and time-dependent movement dynamics. We demonstrate the efficacy of SCMD through empirical validations on the complex molecular system (i.e., single-chain coarse-grained polymers in implicit solvent) for long-duration simulation and dynamical property prediction.