Geometric Spatiotemporal Transformer to Simulate Long-Term Physical Dynamics
Keywords: Equivariance, Spatio-Temporal Transformer, Physical Dynamics
Abstract: Physical dynamics simulation plays a crucial role in various real-world applications. In this paper, we explore the potential of leveraging Transformers by framing the task as autoregressive next-graph prediction based on spatiotemporal graph inputs. To achieve this, we propose Geometric Spatiotemporal Transformers (GSTs), which adopt the expressive encoder-decoder architecture of traditional Transformers. At the core of GSTs are equivariant spatiotemporal blocks that alternate between spatial and temporal modules while preserving E(3) symmetries. Additionally, we introduce the Temporal Difference Graph (TDG), derived from the difference between the last two frames of historical input, to capture global dynamic patterns and mitigate cumulative errors in long-term prediction tasks. Unlike existing Graph Neural Network (GNN) methods, GSTs can process full input sequences of arbitrary lengths to effectively capture long-term context, and address cumulative errors over long-term rollouts thanks to the TDG mechanism. Our method achieves state-of-the-art performance across multiple challenging physical systems at various scales (molecular-, protein-, and macro-level), demonstrating the robust dynamics simulation capabilities.
Primary Area: applications to physical sciences (physics, chemistry, biology, etc.)
Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.
Submission Guidelines: I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2025/AuthorGuide.
Anonymous Url: I certify that there is no URL (e.g., github page) that could be used to find authors’ identity.
No Acknowledgement Section: I certify that there is no acknowledgement section in this submission for double blind review.
Submission Number: 10105
Loading