Geometric Spatiotemporal Transformer to Simulate Long-Term Physical Dynamics

Authors: Zongzhao Li, Jiacheng Cen, Liming Wu, Hao Sun, Hangyu Mao, zhangfuzheng, Di ZHANG, Wenbing Huang
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
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.

Decision: reject
Reviewer scores: 6, 5, 6, 3

Reviews and discussion:
<thread>
creation date:
1729079297894

thread_type:
official_review

<thread_object>
creator: reviewer_ygwj

<content>
summary:
The paper proposes multiple components aimed at improving modeling temporal physics simulations with particular focus on long-horizon rollout performance. At the core of the paper is the Temporal Difference Graph (TDG) which aims to mitigate error accumulation during the rollout by using a history of previous frames together with predicting the delta to the next frame instead of the next frame directly. This is combined with equivariant spatiotemporal transformer blocks as core building block. The proposed method shows improvements over state-of-the-art in small-scale molecular dynamics, motion capture and protein dynamics.

strengths:
- Predicting the delta to the next timeframe is a compelling strategy to mitigate error accumulation. This strategy is also commonly used in diffusion models and the paper shows that this strategy also works well on physical simulation rollouts.
- Using transformers instead of GNNs makes a lot of sense as their scalability is well established and their flexibility (e.g. like using the TDG to incorporate arbitrary many history timesteps) is desireable
- The method shows strong improvements over competitors.
- Clear mathematical and visual description of the method.

weaknesses:
- The paper lacks runtime/flops and parameter count comparisons. Even though the architecture and training procedure can be quite different between methods, a comparison of parameter counts and runtimes is essential to exclude performance gains simply due to investing more compute/parameters into the model/training. 
- Processing the history (especially in the GST-V) could become quite costly, so the runtime necessary for a rollout should be reported.
- The considered models are extremely small (hidden dim 16 with 2 layers according to the Appendix). This limits the interaction capabilities of GNNs and therefore plays into the hand of transformers, which have global interactions in all layers. A comparison of deeper models or a justification why it is not limiting GNNs to be local should be provided (not sure how many message passing steps are common in these setups).
- The problem sizes should be stated in the paper/appendix, it is only stated how many samples are in a dataset, but not how many nodes a single sample has. 
- Autoregressive models are commonly trained with next-step prediction. However due to the time aggregation and history components of the proposed method, a multi-step prediction during training is necessary which greatly limits scalability to larger models, longer rollouts and larger problem sizes (number of nodes)
- The considered datasets are at most 20 rollout timesteps, "long-horizon" seems a bit of a stretch for that.
- Given the similar performance of GST-F and GST-V, it would be interesting to visualize the attention maps to see if GST-V actually uses all of the history or only the recent ones.

questions:
- How do models compare in terms of runtimes/parameters/flops?
- Why are the models so small, are the datasets so small/simple that larger ones run into overfitting? Are larger models infeasible? You are using A100 cards so larger models should not be a problem in terms of compute. Could you elaborate more on the scaling properties of your method (in terms of compute and to larger problem sizes)? 
- Is "temporal position embedding with the sine function" the standard transformer positional embedding with sine/cosine of variable frequencies?
- "We collect all the long-term predictions" makes it seem that long term predictions are emphasized over short term ones, however the MSE weights all predictions equally, right?
- How many timesteps are provided for GST-F?
- In Figure 2, why is $X'_T$ summed with the delta and not $X_T$. Is $X'_T$ just copied from $X_T$ for visualization?

Minor:
- line 72: "inFig" space missing
- line 123: "xpatiotemporal"
- line 264: "the decoder process" should be " the decoder processes"
</content>
</thread_object>

</thread>
<thread>
creation date:
1730065858464

thread_type:
official_review

<thread_object>
creator: reviewer_64hu

<content>
summary:
This paper proposes the Geometric Spatiotemporal Transformer (GST) for predicting physical dynamics in domains such as molecular and motion simulations. Unlike prior works that rely on GNNs, GST leverages the Transformer architecture to handle sequences of arbitrary length, allowing for more flexible long-term predictions. Experimental results show that GST outperforms previous methods, especially in mitigating cumulative errors over extended prediction horizons.

strengths:
- Exploring the efficacy of Transformers over GNNs is an important and interesting research direction.
- The proposed adaptation of the Transformer for physical dynamics is well-motivated and appropriate.
- Experimental results effectively demonstrate the superiority of the proposed method.

weaknesses:
**Limited technical novelty**

The proposed architecture is a straightforward extension of prior work. The equivariant component builds on previous research in geometric deep learning, while the alternating spatio-temporal operations are well-explored in the Video Transformer domain, as seen in works like TimeSformer [1]. In this context, the paper essentially replaces spatial self-attention with an equivariant MPNN.

[1] Bertasius et al. Is Space-Time Attention All You Need for Video Understanding? ICML 2021.

---
**Effectiveness in long-term rediction**

The Temporal Difference Graph (TDG) partially mitigates error accumulation in long-sequence predictions, similar to common tricks such as using advantage functions to improve over value functions. However, this approach does not fundamentally address the long-sequence issue, as the accumulation of temporal differences also grows over time.

The experimental results do not convincingly demonstrate GST’s specific advantage for long-sequence predictions. In Table 2, GST consistently outperforms prior methods across all horizons but does not exhibit a marked improvement at longer horizons. To substantiate the claim of efficacy on long sequences, the paper should provide evidence of stable performance over extended horizons, while prior methods should show significant performance degradation in these cases.

questions:
N/A
</content>
</thread_object>

</thread>
<thread>
creation date:
1730572987900

thread_type:
official_review

<thread_object>
creator: reviewer_v53b

<content>
summary:
The paper introduces a Geometric Spatiotemporal Transformer (GST) designed for long-term simulation of physical dynamics using Transformers in a novel spatiotemporal context. Key innovations include E(3)-equivariant spatiotemporal blocks and a Temporal Difference Graph (TDG) for cumulative error reduction. The model is evaluated across molecular, protein, and motion datasets, aiming to outperform existing GNN-based methods.

strengths:
- This paper proposes a novel dynamical model that can handle spatiotemporal sequences of variable length.
- The authors provide theoretical proof and show comprehensive experimental comparisons.

weaknesses:
- In addition to EGNN, there are other newer equivariant graph neural network architectures. Why not use them or conduct comparative experiments? EGNN has been proven not to be truly equivariant.
- In the temporal module, an equivariant self-attention mechanism is used to model each node. Are these nodes shared parameters? If not, the question of how to deal with the non-constant number of nodes in the data set should be answered to reflect the applicability of the model.
- Where is causal attention strategy reflected, and how is it different from traditional attention mechanism?
- Transformer-based models normally have high time complexity. The authors should show the efficiency comparison with other methods.
- Predicting the differences instead of directly predicting the value of the next frame has long been used in methods such as GNS[1] and MGN[2]. Why not compare?
- The detailed architectural differences between GST-V and GST-F should be explained. GST can be extended to variable lengths based on attention, so why can't ESTAG?
- Why can $\boldsymbol{X}_{\delta}$ reduce the cumulative error in long-term predictions? This should only be a short-term bias term. What is the theoretical basis?

[1] A Sanchez-Gonzalez et al. 'Learning to simulate complex physics with graph network.' ICML2020.

[2] T Pfaff et al. 'Learning Mesh-Based Simulation with Graph Networks.' ICLR2021.

questions:
Please address the questions in the Weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730645731965

thread_type:
official_review

<thread_object>
creator: reviewer_daxk

<content>
summary:
This paper presents a Transformer-based approach Geometric Spatiotemporal Transformers for simulating long-term physical dynamics. To improve the handling of cumulative errors over extended time horizons, this paper frames the task as autoregressive next-graph prediction, integrating E3 equivariant spatiotemporal blocks, and a TDG. It is evaluated across various scales, including molecular, protein, and macro-level physical systems, showing improvements in long-term predictive accuracy over existing methods.

strengths:
1. An Innovative Use of Transformers is presented. Integrating Transformer architectures with spatiotemporal graph structures and E(3) equivariance introduces a novel approach for handling long-term physical dynamics, marking a unique departure from traditional GNN-based methods.
2. TDG shows promise in addressing cumulative errors in long-rollout predictions in physical simulation tasks.
3. Experiments offer some insights into the contributions of individual components, including equivariance and TDG which helps evaluate the effectiveness of the architecture. The model’s performance is assessed across various physical systems at molecular, protein, and macro levels, demonstrating GST’s potential adaptability to different scales of physical dynamics.

weaknesses:
1. The experiments should include at least one additional dataset to better demonstrate generalizability, especially to popular datasets where multi-frame prediction methods may underperform. This would strengthen the paper's claims of achieving SOTA performance.
2. Many of the reported improvements over existing models are marginal and are often limited to specific datasets or rollout lengths. The author should clearly frame these gains to avoid overstating the model’s broader applicability.
3. The architecture's complexity, with components like equivariant modules and TDG, presents a reproducibility challenge. Offering an easy-to-use module, similar to Transformers, could significantly improve accessibility and replicability for other researchers.

questions:
1. Can the TDG framework be adapted or tuned for other symmetry types beyond E(3) when applied to different dynamic systems?
2. Could the authors provide a clear, modular implementation that enables easy usage and integration later?
</content>
</thread_object>

</thread>
<thread>
creation date:
1733313476424

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response

comment:
We sincerely thank all reviewers and ACs for their time and efforts on reviewing the paper. We are very glad that the reviewers recognized the problems we studied, and the models we built, and their comments really gave us a lot of inspiration and enlightenment.

In this paper, we explore the potential of leveraging Transformers by framing the physical dynamics simulation as autoregressive next-graph prediction based on spatiotemporal graph inputs. Our proposed Geometric Spatiotemporal Transformers (GSTs) processes full input sequences of arbitrary lengths to effectively capture long-term context, and address cumulative errors over long-term rollouts. As the first work to explore the autoregressive next-graph prediction paradigm on tasks (MD17, Protein, Motion Capture) within the physical dynamics domain, we hope this study will draw attention and foster the development of this paradigm across the field. Any feedback provided during the rebuttal process will be thoroughly considered and integrated into the final version.   

To address the reviewers' concerns, we have added additional experiments and figures as follows.

**Table A** in the response to Reviewer DaxK reports the performance improvement on MD17 and Protein Dynamics.

**Table B** in the response to Reviewer DaxK reports the performance improvement on Motion Capture.

**Table A** in the response to Reviewer V53b compares inference time between GST with other methods.

**Table B** in the response to Reviewer V53b shows the performance of GST versus GNS on MD17.

**Table C** in the response to Reviewer V53b shows the performance of GST when using second-order differences on MD17.

**Table A** in the response to Reviewer 64HU shows the average per-frame prediction MSE difference between our method and the previous SOTA method.

**Table A** in the response to Reviewer YGWJ compares inference runtimes and parameter counts of all models for the Aspirin molecule.

**Table B** in the response to Reviewer YGWJ reports the model's performance on the Aspirin molecule for Layer=2,3,4,5,6.

**Figure 6** in Appendix D illustrates the visualization of the attention map for GST-V.

**The figure** in https://anonymous.4open.science/r/GST-DFC2/malonaldehyde.png shows the results of the previous SOTA method ESTAG and GST on 200 rollout steps on MD17.
</content>
</thread_object>

</thread>
<thread>
creation date:
1733313572958

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Rebuttal Acknowledgment

comment:
Dear Reviewers and Area Chairs,

We would like to extend our heartfelt gratitude for your dedicated efforts, insightful feedback, and constructive suggestions. Your nice comments and thoughtful suggestions are truly commendable.

Through our discussions and the reviewers' responses, we believe we have addressed the major concerns raised by all reviewers. Among the rebuttal, 2 out of 4 reviewers (Reviewer V53b and Reviewer 64HU) raised their ratings. Specifically, Reviewer V53b raised from 3 to 5, and Reviewer 64HU raised from 5 to 6. Reviewer YGWJ intended to raise the score from 3 to 4 but kept it unchanged due to the lack of this option. This outcome has greatly benefited and encouraged us, and we would like to thank all of you for your valuable support!

To address Reviewer YGWJ’s concerns, we provided detailed responses to each question during the rebuttal phase, supported by experimental results. In response to our initial reply, Reviewer YGWJ acknowledged that some issues had been resolved but raised three remaining concerns. For the "long-horizon" issue, we provided more detailed explanations and presented the model’s performance under a 200-step rollout setting as requested. Regarding the "scalability of this approach," we further clarified the model’s potential to handle larger problem sizes from the perspective of time complexity. As for the suggestion that "the answers to questions from the review should be included in the paper," we incorporated these points into the revised manuscript. However, due to policy restrictions, we are unable to upload the updated PDF at this stage. By the end of the rebuttal phase, we did not receive further feedback from Reviewer YGWJ on our latest responses.

We acknowledge that some expressions in the initial draft might have been unclear, but we have provided detailed explanations. We hope that our detailed responses will help the AC and reviewers gain a clearer understanding of our paper and resolve any potential misunderstanding.

Once again, we express our sincere appreciation for your time and dedication in reviewing our work. Your insightful input has significantly contributed to the refinement of our manuscript. Thank you!

Best regards,

The Authors
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper introduces a Transformer-based model to simulate long-term physical dynamics of graph-structured data. The main contributions of the paper are (1) introducing E(3)-equivariant modules into the Transformer blocks to respect underlying symmetries in three-dimensional space and (2) introducing temporal difference graph to mitigate error accumulation in long-horizon rollouts. 

The reviewers were divided. Two reviewers ended with a slightly positive stance, one with a slightly negative stance, and one firmly on the reject side. The primary concerns from the reviewers were about (1) incremental novelty, since the proposed method is primarily an incremental combination of prior works (e.g., equivariance from prior work and difference predictions similar to other simulation approaches), (2) limited justifications on long-horizon rollout, since the experiments are mainly conducted on small scales (e.g., 20 steps), and (3) scalability issues. While authors response partly addressed some of the important concerns, some reviewers were not fully convinced with the responses, especially regarding the weakness (2) and (3).  

After reading the paper, reviews, and discussions carefully, AC believes that the primary contributions of the paper—adaptation of Transformers with equivariance and a TDG mechanism—can be a promising direction. However, the AC also agrees with the reviewers that the results on long-horizon rollout presented in the paper are still insufficient to support the authors’ claim, since they are mostly constrained on moderately short-term frame (<20 steps). While the preliminary long-term prediction results presented in the rebuttal are encouraging, it remains insufficient to assess its long-horizon rollout quality and generalizability over other datasets. Given that these limitations do not overweight the limited technical novelty, AC believes that this paper does require significant work before it can be published.