Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Equivariant Spatio-Temporal Attentive Graph Networks to Simulate Physical Dynamics

Liming Wu, Zhichao Hou, Jirui Yuan, Yu Rong, Wenbing Huang

TL;DR

ESTAG addresses non-Markovian dynamics in physical systems by reformulating dynamics as a spatio-temporal prediction task and enforcing $E(3)$-equivariance. It introduces Equivariant Discrete Fourier Transform (EDFT) to extract frequency features, and combines Equivariant Spatial Module (ESM) with Equivariant Temporal Module (ETM) to perform iterative spatial and temporal message passing. The model demonstrates superior accuracy across molecular, protein, and macro-scale datasets, with ablations confirming the essential roles of EDFT, attention, and equivariance. This approach provides a principled, symmetry-preserving framework for simulating complex dynamics with unseen environmental factors, offering potential impact in molecular dynamics, protein folding, and robotics. Future work could integrate energy-conservation priors and multi-scale GNNs to handle larger systems and longer horizons.

Abstract

Learning to represent and simulate the dynamics of physical systems is a crucial yet challenging task. Existing equivariant Graph Neural Network (GNN) based methods have encapsulated the symmetry of physics, \emph{e.g.}, translations, rotations, etc, leading to better generalization ability. Nevertheless, their frame-to-frame formulation of the task overlooks the non-Markov property mainly incurred by unobserved dynamics in the environment. In this paper, we reformulate dynamics simulation as a spatio-temporal prediction task, by employing the trajectory in the past period to recover the Non-Markovian interactions. We propose Equivariant Spatio-Temporal Attentive Graph Networks (ESTAG), an equivariant version of spatio-temporal GNNs, to fulfill our purpose. At its core, we design a novel Equivariant Discrete Fourier Transform (EDFT) to extract periodic patterns from the history frames, and then construct an Equivariant Spatial Module (ESM) to accomplish spatial message passing, and an Equivariant Temporal Module (ETM) with the forward attention and equivariant pooling mechanisms to aggregate temporal message. We evaluate our model on three real datasets corresponding to the molecular-, protein- and macro-level. Experimental results verify the effectiveness of ESTAG compared to typical spatio-temporal GNNs and equivariant GNNs.

Equivariant Spatio-Temporal Attentive Graph Networks to Simulate Physical Dynamics

TL;DR

ESTAG addresses non-Markovian dynamics in physical systems by reformulating dynamics as a spatio-temporal prediction task and enforcing -equivariance. It introduces Equivariant Discrete Fourier Transform (EDFT) to extract frequency features, and combines Equivariant Spatial Module (ESM) with Equivariant Temporal Module (ETM) to perform iterative spatial and temporal message passing. The model demonstrates superior accuracy across molecular, protein, and macro-scale datasets, with ablations confirming the essential roles of EDFT, attention, and equivariance. This approach provides a principled, symmetry-preserving framework for simulating complex dynamics with unseen environmental factors, offering potential impact in molecular dynamics, protein folding, and robotics. Future work could integrate energy-conservation priors and multi-scale GNNs to handle larger systems and longer horizons.

Abstract

Learning to represent and simulate the dynamics of physical systems is a crucial yet challenging task. Existing equivariant Graph Neural Network (GNN) based methods have encapsulated the symmetry of physics, \emph{e.g.}, translations, rotations, etc, leading to better generalization ability. Nevertheless, their frame-to-frame formulation of the task overlooks the non-Markov property mainly incurred by unobserved dynamics in the environment. In this paper, we reformulate dynamics simulation as a spatio-temporal prediction task, by employing the trajectory in the past period to recover the Non-Markovian interactions. We propose Equivariant Spatio-Temporal Attentive Graph Networks (ESTAG), an equivariant version of spatio-temporal GNNs, to fulfill our purpose. At its core, we design a novel Equivariant Discrete Fourier Transform (EDFT) to extract periodic patterns from the history frames, and then construct an Equivariant Spatial Module (ESM) to accomplish spatial message passing, and an Equivariant Temporal Module (ETM) with the forward attention and equivariant pooling mechanisms to aggregate temporal message. We evaluate our model on three real datasets corresponding to the molecular-, protein- and macro-level. Experimental results verify the effectiveness of ESTAG compared to typical spatio-temporal GNNs and equivariant GNNs.
Paper Structure (29 sections, 2 theorems, 17 equations, 12 figures, 9 tables, 1 algorithm)

This paper contains 29 sections, 2 theorems, 17 equations, 12 figures, 9 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Work
  3. Notations and Task Definition
  4. Our approach: ESTAG
  5. Equivariant Discrete Fourier Transform (EDFT)
  6. Equivariant Spatial Module (ESM)
  7. Equivariant Temporal Module (ETM)
  8. Experiments
  9. Molecular-level: MD17
  10. Protein-level: Protein Dynamics
  11. Macro-level: Motion Capture
  12. Ablation Studies
  13. Conclusion
  14. Acknowledgement
  15. Proof of ESTAG's Equivariance
  16. ...and 14 more sections

Key Result

Theorem 4.1

We denote ESTAG as $\vec{{\bm{X}}}(T) =\phi\left(\{({\bm{H}}(t), g\cdot\vec{{\bm{X}}}(t),{\bm{A}})\}_{t=0}^{T-1}\right)$, then $\phi$ is E($3$)-equivariant. Proof. See Appendix sec:theorem proof.

Figures (12)

  • Figure 1: Comparison of the problem setting between previous methods and our paper. Here, we choose the dynamics of the Aspirin molecular with time lag as 1 for illustration.
  • Figure 2: Schematic overview of ESTAG. After inputting historical graph trajectories ${\mathcal{G}}_0,...,{\mathcal{G}}_{T-1}$, Equivariant Discrete Fourier Transform (EDFT) extracts equivariant frequency features $\vec{{\bm{f}}}$ from the trajectory. We process them into the invariant node-wise feature ${\bm{c}}$ and adjacency matrix ${\bm{A}}$ to be adopted for the next stage. Then we stack Equivariant Spatial Module (ESM) and Equivariant Temporal Module (ETM) alternatively for $L$ times to explore spatial and temporal dependencies. After the equivariant temporal pooling layer, we obtain the estimated position $\vec{{\bm{x}}}^*(T)$.
  • Figure 3: Visualization of cross-correlation ${\bm{A}}$ on Aspirin. EDFT can not only identify strongly-connected nodes (e.g. Node 8 and Node 11), but also discover latent relationship between two nodes which are disconnected yet may have similar structures or functions (e.g. Node 8 and Node 10).
  • Figure 4: PyMol visualization of the predicted molecules by our ESTAG and ST_EGNN, where the MSE ($\times 10^{-3}$) with respect to the ground truth is also shown. As expected, the predicted instances by ESTAG exhibit much smaller MSE than ST_EGNN, although the difference is not easy to visualize in some cases. For those obviously mispredicted regions of ST_EGNN, we highlight them with red rectangles. It is observed that ST_EGNN occasionally outputs isolated atoms, which could be caused by violation of the bond length tolerance in PyMol.
  • Figure 5: The rollout-MSE curves on 8 molecules in MD17. Our model generally achieves the lowest MSE.
  • ...and 7 more figures

Theorems & Definitions (3)

  • Theorem 4.1
  • Theorem A.1
  • proof