Physical Dynamics as Next Geometric Graph Prediction

Zongzhao Li, Jiacheng Cen, Liming Wu, Hao Sun, Ruihua Song, Hangyu Mao, zhangfuzheng, Di ZHANG, Wenbing Huang

Keywords: Equivariance, Spatio-Temporal Transformer

Abstract: Physical dynamics simulation serves as a foundational component in scientific computing and AI applications. This paper presents a novel approach that redefines the problem as autoregressive prediction of spatiotemporal graph sequences. Built upon the expressivity of Transformer, we propose an Equivariant Spatiotemporal Transformer (EST), extending conventional Transformers with specialized equivariant spatiotemporal blocks. These blocks systematically alternate between spatial and temporal modules, rigorously maintaining E(3) symmetries throughout the process. Moreover, the design incorporates a novel Temporal Difference Graph (TDG) module derived from frame-wise variations, effectively modeling global dynamics and addressing cumulative errors in autoregressive predictions. Unlike traditional graph neural networks, our EST can process variable-length historical sequences and mitigate the persistent challenge of error accumulation in autoregressive processes. Comprehensive evaluations across multiscale physical systems (molecular-, protein-, and macroscopic-scale) demonstrate that our method achieves state-of-the-art performance, thereby showcasing its robust and versatile dynamics simulation capabilities.

Primary Area: learning on graphs and other geometries & topologies

Submission Number: 9097