[Paper Review] Equivariant Spatio-Temporal Attentive Graph Networks to Simulate Physical Dynamics

Liming Wu, Zhichao Hou+3 more|arXiv (Cornell University)|May 21, 2024

Advanced Graph Neural Networks4 citations

TL;DR

This paper proposes ESTAG, an equivariant spatio-temporal graph network that models non-Markovian physical dynamics by leveraging past trajectory history. It integrates an Equivariant Discrete Fourier Transform (EDFT) for periodic pattern extraction, an Equivariant Spatial Module (ESM), and an Equivariant Temporal Module (ETM) with forward attention and equivariant pooling, achieving state-of-the-art performance across molecular, protein, and macro-scale dynamics benchmarks.

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.

Motivation & Objective

• Address the limitation of Markovian assumptions in existing physical dynamics simulation methods, which fail to capture unobserved environmental interactions.
• Reformulate dynamics simulation as a spatio-temporal prediction task using historical trajectory data to model non-Markovian dependencies.
• Develop an equivariant architecture that preserves Euclidean symmetries (rotation, translation) while modeling temporal and spatial dependencies.
• Integrate periodic pattern modeling via a novel Equivariant Discrete Fourier Transform (EDFT) to capture recurring dynamics in physical systems.
• Demonstrate the superiority of combining temporal memory and equivariance over standard spatio-temporal GNNs and non-equivariant models.

Proposed method

• Propose a novel Equivariant Discrete Fourier Transform (EDFT) to extract frequency-domain features from historical node trajectories while preserving E(3) equivariance.
• Design an Equivariant Spatial Module (ESM) that performs message passing on spatial graphs using equivariant updates to node coordinates and features.
• Introduce an Equivariant Temporal Module (ETM) with forward attention and equivariant pooling to aggregate temporal messages across frames while maintaining geometric symmetry.
• Construct spatio-temporal graphs from historical system states to model long-range temporal dependencies beyond the frame-to-frame assumption.
• Ensure all learned representations and updates are equivariant under translations, rotations, and reflections via careful design of message passing and aggregation layers.
• Use a learnable attention mechanism in ETM to dynamically weight the importance of past frames, enabling adaptive temporal modeling.

Experimental results

Research questions

• RQ1 Can modeling physical dynamics as a spatio-temporal prediction task, rather than a frame-to-frame forecast, better capture non-Markovian behaviors caused by unobserved environmental interactions?
• RQ2 How can we design an equivariant spatio-temporal GNN that preserves E(3) symmetry while effectively modeling long-range temporal dependencies?
• RQ3 What is the contribution of periodic pattern modeling via EDFT to the performance of physical dynamics simulation?
• RQ4 To what extent does combining temporal memory with geometric equivariance improve generalization and accuracy compared to standard STGNNs or equivariant GNNs?
• RQ5 How do individual components—EDFT, attention, equivariance, temporal pooling—contribute to the overall model performance?

Key findings

• ESTAG achieves state-of-the-art performance on MD17, Protein, and Macro-level datasets, with a mean MSE of 0.063 on Aspirin and 0.003 on Benzene, significantly outperforming non-equivariant and non-temporal baselines.
• Abllation studies show that removing EDFT increases MSE by up to 25% on average, demonstrating its critical role in modeling spatial periodicity.
• Removing the attention mechanism in ETM increases MSE by 38% on average, indicating its importance in capturing dynamic temporal interactions.
• The non-equivariant version of ESTAG shows a 10x increase in MSE on average, proving that geometric equivariance is essential for robust physical simulation.
• Replacing equivariant temporal pooling with learnable coefficients leads to a 35% increase in MSE, confirming the advantage of structured temporal aggregation.
• With L=2 layers, ESTAG achieves optimal trade-off between accuracy and efficiency on the Ethanol dataset, with an MSE of 0.990×10⁻⁴.