Non-stationary Equivariant Graph Neural Networks for Physical Dynamics Simulation

To enhance the generalization ability of graph neural networks (GNNs) in learning and simulating physical dynamics, a series of equivariant GNNs have been developed to incorporate the symmetric inductive bias. However, the existing methods do not consider the non-stationarity nature of physical dynamics, where the joint distribution changes over time. Moreover, previous approaches for modeling nonstationary time series typically involve normalizing the data, which disrupts the symmetric assumption inherent in physical dynamics. To model the non-stationary physical dynamics while preserving the symmetric inductive bias, we introduce a Non-Stationary Equivariant Graph Neural Network (NS-EGNN) to capture the non-stationarity in physical dynamics while preserving the symmetric property of the model. Specifically, NS-EGNN employs Fourier Transform on segments of physical dynamics to extract time-varying frequency information from the trajectories. It then uses the first and second-order differences to mitigate non-stationarity, followed by pooling for future predictions. Through capturing varying frequency characteristics and alleviate the linear and quadric trend in the raw physical dynamics, NS-EGNN better models the temporal dependencies in the physical dynamics. NS-EGNN has been applied on various types of physical dynamics, including molecular, motion and protein dynamics, and consistently surpasses the existing state-of-the-art algorithms, underscoring its effectiveness. The implementation of NS-EGNN is available at https://github.com/MaojiWEN/NS-EGNN.