Geometric Spatiotemporal Transformer to Simulate Long-Term Physical Dynamics

Your detailed review has been immensely helpful! We have systematically addressed each of your comments, and our responses are presented below.

W1 & Q1: The paper lacks runtime/flops and parameter count comparisons.

A1: Nice suggestion! In Table A, we present the full inference runtimes of our model and other baseline models on the MD17 dataset for the Aspirin molecule under the rollout=10 setting, along with the parameter counts of all models. To demonstrate that the performance improvement of our model is not solely due to an increase in parameter counts, we also report the performance of our model under a comparable parameter count with ESTAG in the ablation study (Table 4, line 6 Shared Parameters) of the original paper.

Table A: Inference runtimes and parameter counts of all models for the Aspirin molecule.

W2: Processing the history (especially in the GST-V) could become quite costly, so the runtime necessary for a rollout should be reported.

A2: Thank you for your suggestion! We report the corresponding results in Table A. As you correctly pointed out, GST-V requires longer inference time compared to other methods due to processing the full historical sequence. However, this additional computation brings significantly improved performance.

W3: 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).

A3: Nice suggestion! There may be some misunderstanding here. In our paper, we use the EGNN architecture to model geometric information in the spatial dimension, which is consistent with other baseline models (e.g., ST_GNN, ST_EGNN, ESTAG) because the 3D structural topological data of the input is well-suited for GNN modeling. Additionally, since the number of input time frames can theoretically be very large, we use a self-attention mechanism in the temporal dimension to model the relationships between the same node across different frames. Thus, our model design does not limit the interaction capabilities of GNNs but rather leverages the distinct strengths of GNNs and Transformer architectures to model different dimensions effectively.

W4: 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.

A4: Thank you for your suggestion! We have added the number of nodes contained in a single sample for each dataset in Appendix B of the revised paper.

W5: 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).

A5: Thank you for your feedback! As you pointed out, extending GST-V to longer rollouts is indeed a significant challenge. However, this issue is common to all autoregressive models that consider historical components and remains difficult to solve. In this paper, we conducted a preliminary exploration in this direction, similar to the recent attempt by [1] on time series tasks. For scaling to larger models, we believe this is currently constrained by the simplicity of the datasets, where larger models tend to overfit. Regarding scaling to larger problem sizes (number of nodes), we consider this feasible, as it mainly depends on the design of the EGNN for modeling spatial geometric relationships. For instance, setting a distance threshold to limit the number of graph edges can effectively reduce model complexity.

W6: The considered datasets are at most 20 rollout timesteps, "long-horizon" seems a bit of a stretch for that.

A6: Thank you for your detailed observation! Although the paper mentions a maximum of 20 rollout timesteps, we provide a description of the Time Lag for sampled timestamps in Appendix B. In the MD17 dataset, with a Time Lag of 10, 20 rollout timesteps actually correspond to an input sequence of 200 frames. We believe this can be considered a relatively long horizon.

W7: 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.

A7: Thank you for your suggestion! In Figure 6 of Appendix D, we present the visualization of the attention map for GST-V. The visualization shows that most timesteps attend to a significant portion of the historical frames. This further validates that GST-V effectively captures and integrates historical input information to enhance prediction performance.

Q2: Could you elaborate more on the scaling properties of your method (in terms of compute and to larger problem sizes)?

Answer to Q2: Nice suggestion! In the MD17 and Motion Capture datasets, the number of nodes per sample is relatively small (a maximum of 13 nodes in MD17 and 31 nodes in Motion Capture), making larger models prone to overfitting. Additionally, there is a lack of datasets in the field describing more challenging and larger molecular dynamics problems, which we plan to explore in future research. On the Protein Dynamics dataset, where the number of nodes is 5325, the results reflect the scaling properties of our model for larger problem sizes.

To more intuitively demonstrate the overfitting issue with larger model sizes, we provide the model's performance on the Aspirin molecule for L=2,3,4,5,6 in Table B. The experimental results suggest that increasing L does not improve the model's performance, indicating a potential overfitting phenomenon.

Table B: The model's performance on the Aspirin molecule for L=2,3,4,5,6.

Q3: Is "temporal position embedding with the sine function" the standard transformer positional embedding with sine/cosine of variable frequencies?

Answer to Q3: Yes, we use the sine/cosine positional embedding form from the standard Transformer architecture for the temporal position embedding.

Q4: "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?

Answer to Q4: Thank you for your detailed observation! The statement “We collect all the long-term predictions” means that we gather results from multiple rollouts and compute the MSE loss collectively. During loss computation, each frame’s loss is compared to the corresponding ground truth, and all frames are treated equally with the same weight.

Q5: How many timesteps are provided for GST-F?

Answer to Q5: We provide 10 timesteps, the same as other baseline methods.

Q6: 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?

Answer to Q6: Thank you for your detailed observation! Here, $X'\_{T}$ represents the model output for the coordinates of frame $T$, and similarly, $X'\_{\delta}$ represents the coordinate information from the TDG output. We predict the coordinates of frame $T+1$ by summing $X'\_{T}$ and $X'\_{\delta}$.

Q7: The typos

Answer to Q7: Thank you for pointing out the typos! We have corrected them in the revised version of the paper.

[1] Liu Y et al. AutoTimes: Autoregressive Time Series Forecasters via Large Language Models, NIPS 2024.

Dear Reviewer YGWJ,

We extend our sincerest gratitude for your time and valuable feedback in reviewing our paper! We have meticulously responded to all of your questions and concerns, and we would like to know if you are satisfied with our responses and willing to reconsider your evaluation and score for our paper.

As the rebuttal deadline is approaching, we sincerely hope to receive your feedback at your earliest convenience.

Best regards,

The Authors

Dear Reviewer YGWJ,

With only a few hours left until the rebuttal period ends, we would like to ask if you are satisfied with our latest responses? We would greatly appreciate it if you could reassess our paper and consider raising your score!

Best regards,

The Authors

We sincerely appreciate the time and effort you invested in reviewing our paper! We have carefully considered and responded to each of your remarks.

W1: Limited technical novelty

A1: Thank you for your feedback! First, we are the first to adopt standard Transformer architecture in the physical dynamics domain. Similar to ViT [1], our exploration of this new scenario is itself a contribution. Second, we emphasize that transferring such architectures to this domain is non-trivial, as spatiotemporal tasks in videos and 3D geometric graph processing differ significantly:

• Data format: Video tasks take video sequences as input, whereas our input, after feature extraction, becomes 3D structural data with topological relationships.
• Equivariance: Video-based models (e.g., TimeSformer) do not consider equivariance, which is a key property in our tasks. Our approach is not a simple alternation of spatial EGNN modules and temporal attention modules but a tailored design for equivariance. Specifically, for each node, we ensure both intra-frame atomic interaction equivariance and temporal consistency, maintaining spatial and temporal equivariance. In our paper, we introduce the Temporal Difference Graph (TDG), integrating attention mechanisms in the encoder and decoder with equivariant forms. We also ensure TDG equivariance through data decentralization. Moreover, the TDG design effectively reduces accumulated errors, further improving the model's long-term prediction performance.

W2: Effectiveness in long-term rediction

A2: Thank you for your constructive suggestions! As you pointed out, introducing TDG can only partially mitigate the impact of accumulated errors. However, accumulated error remains a significant challenge. As the rollout steps increase, predictions inevitably result in progressively larger errors, which is a difficulty faced by nearly all models. Therefore, our work simply represents initial exploration in this direction.

We also acknowledge that we failed to clearly demonstrate the significant advantages of our model in long-sequence prediction tasks, and we sincerely apologize for this. Below are further details:

• The vertical axis in Figure 3 of the paper is log-transformed. When using raw data, the differences between curves would be more pronounced.
• We reanalyzed Table 2 and computed the average per-frame prediction MSE difference (normalized by rollout steps) between our method and the previous SOTA method ESTAG, as shown in Table A. The results indicate that as rollout steps increase, the gap between our method and ESTAG grows, demonstrating the superior performance of our model in long-range prediction tasks.

Table A: The average per-frame prediction MSE difference between our method and the previous SOTA method.

[1] Dosovitskiy A et al. An image is worth 16x16 words: Transformers for image recognition at scale, ICLR 2021.

Dear Reviewer 64HU,

We extend our sincerest gratitude for your time and valuable feedback in reviewing our paper! We have meticulously responded to all of your questions and concerns, and we would like to know if you are satisfied with our responses and willing to reconsider your evaluation and score for our paper.

As the rebuttal deadline is approaching, we sincerely hope to receive your feedback at your earliest convenience.

Best regards,

The Authors

The update formula for the former is (where $\vec{\boldsymbol{x}}_{t,c}^{e,l}$ denotes the centroid of the geometric graph):

$\vec{\boldsymbol{x}}\_{t,i}^{(e,l+1)}=\vec{\boldsymbol{x}}\_{t,i}^{(e,l)}+\frac{1}{|\mathcal{N}(i)|}\sum\_{j\in\mathcal{N}}\left[\varphi\_x\left(\boldsymbol{m}\_{t,ij}^{e,l}\right)\cdot\left(\vec{\boldsymbol{x}}\_{t,i}^{e,l}-\vec{\boldsymbol{x}}\_{t,j}^{e,l}\right)+\varphi\_{\text{cross}}\left(\boldsymbol{m}\_{t,ij}^{e,l}\right)\cdot\left(\left(\vec{\boldsymbol{x}}\_{t,i}^{e,l}-\vec{\boldsymbol{x}}\_{t,c}^{e,l}\right)\times\left(\vec{\boldsymbol{x}}\_{t,j}^{e,l}-\vec{\boldsymbol{x}}\_{t,c}^{e,l}\right)\right)\right]$

The update formula for the latter is:

$\vec{\boldsymbol{x}}\_{t,i}^{(e,l+1)}=\vec{\boldsymbol{x}}\_{t,i}^{(e,l)}+\frac{1}{|\mathcal{N}(i)|}\sum\_{j\in\mathcal{N}}\left[\varphi\_x\left(\boldsymbol{m}\_{t,ij}^{e,l},\vec{\boldsymbol{g}}^\top\left(\vec{\boldsymbol{x}}\_{t,i}^{e,l}-\vec{\boldsymbol{x}}\_{t,j}^{e,l}\right)\right)\cdot\left(\vec{\boldsymbol{x}}\_{t,i}^{e,l}-\vec{\boldsymbol{x}}\_{t,j}^{e,l}\right)+\varphi\_{g}\left(\boldsymbol{m}\_{t,ij}^{e,l},\vec{\boldsymbol{g}}^\top\left(\vec{\boldsymbol{x}}\_{t,i}^{e,l}-\vec{\boldsymbol{x}}\_{t,j}^{e,l}\right)\right)\cdot\vec{\boldsymbol{g}}\right]$

[1] Han J et al. Equivariant Graph Hierarchy-Based Neural Networks, NIPS 2022.

[2] Wu L et al. Equivariant Spatio-Temporal Attentive Graph Networks to Simulate Physical Dynamics, NIPS 2023.

[3] Xu C et al. EqMotion: Equivariant Multi-agent Motion Prediction with Invariant Interaction Reasoning, ICCV 2023.

[4] Xu M et al. Equivariant Graph Neural Operator for Modeling 3D Dynamics, ICML 2024.

[5] Schneuing A et al. Structure-based drug design with equivariant diffusion models, arXiv 2022.

[6] Han J et al. Learning physical dynamics with subequivariant graph neural networks, NIPS 2022.

Thank you for taking the time to review our paper and providing constructive suggestions! We have addressed each of your questions individually, and you can find the corresponding responses below.

W1: 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.

A1: Thank you for your suggestion! The three datasets listed in our paper are widely recognized and popular in this field. Specifically, MD17 and Motion Capture are used as primary experimental datasets in [2, 3, 4], while the Adk equilibrium trajectory dataset is also used as a primary dataset in [1, 2, 4]. Therefore, our selection covers the essential and commonly used datasets in the current research domain as comprehensively as possible.

W2: 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.

A2: Thank you for your feedback! The performance improvement of our model over others is not limited to a single dataset or specific rollout lengths. Moreover, the improvement is both significant and consistent. Table A and Table B summarizes the performance improvement of our model compared to previous SOTA methods across MD17, Protein, and Motion Capture datasets with different rollout lengths (note that the MD17 dataset includes 8 molecules, so we report the average improvement). From Table A and Table B, it is clear that our model demonstrates significant performance improvement across all datasets and all rollout lengths.

Table A: Performance improvement on MD17 and Protein Dynamics

Table B: Performance improvement on Motion Capture

W3 & Q2: 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.

A3: Thank you for your suggestion! Based on your advice, we plan to implement the GST module as a highly modular function and design corresponding interfaces. When adapting to different tasks, users will only need to set the hyperparameters and adjust the input format according to our documentation. We will release the code, relevant function interfaces, and documentation once the paper is accepted.

Q1: Can the TDG framework be adapted or tuned for other symmetry types beyond E(3) when applied to different dynamic systems?

Answer to Q1: Thank you for your feedback! Our TDG network is a highly compatible framework that can be modified to handle SE(3)-equivariance/O$\_g$(3)-equivariance (the latter represents partial symmetry under external forces such as gravity). These two operations can be implemented by following approaches similar to DiffSBDD [5] and SGNN [6], respectively, with adjustments to the coordinate update formula in Eq. (2).

Thank you for your insightful review and constructive comments! Below, we provide detailed responses to each of your queries.

Q1: 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.

A1: Thank you for your suggestion! To the best of our knowledge, EGNN satisfies equivariance and has been widely studied, with numerous works published in top conferences and journals [1,2,3,4]. We chose EGNN as our backbone for the following reasons:

• It is a common practice to extend new architectures based on EGNN, as seen in works like EDM [1] and EGNO [2].
• While stronger equivariant graph neural networks might achieve better results, it would be difficult to determine whether the performance improvement is due to the equivariant network or our GST network.
• Models based on high-degree steerable features (e.g., TFN [5], Equiformer [6]) incur prohibitive computational costs and are not acceptable in terms of time complexity. Additionally, we assume you are referring to the lack of complete expressiveness in EGNN (rather than questioning its equivariance), which is correct. However, we must emphasize that completeness does not necessarily correlate with practical performance. For instance, although TFN has been proven to be complete, it performs worse than EGNN on large-scale dynamics tasks.

Q2: 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.

A2: Yes, each node shares parameters. Therefore, our model is applicable to arbitrary node inputs and scenarios where the number of input nodes is non-constant.

Q3: Where is causal attention strategy reflected, and how is it different from traditional attention mechanism?

A3: Thank you for your question! The causal attention strategy is incorporated into our temporal module, specifically in the Equivariant Temporal Self-Attention module. Since time series have causality and require maintaining physical rationality, predictions for the current frame can only depend on past frames and not on future frames. To implement this, we set a mask in the module to enable causal attention, ensuring that the current frame attends only to past frames and not future frames. In contrast, traditional self-attention mechanisms lack this masking feature, allowing each position in the input sequence to attend to all other positions. This causal attention technique has been widely adopted in various natural language models [7, 8].

Q4: Transformer-based models normally have high time complexity. The authors should show the efficiency comparison with other methods.

A4: Thank you for your suggestion! We have detailed the inference time comparison between our model and others in the Table A below. In practice, the time complexity of the Transformer is $O(T^2)$ (where T is the length of the timestamps). However, since $T$ ranges between 10 and 20 in our settings, the additional computational overhead is minimal. Moreover, given the performance improvement of our model, this additional time cost is acceptable.

Table A: Inference time comparison between our model and others.

Q5: Predicting the differences instead of directly predicting the value of the next frame has long been used in methods such as GNS and MGN. Why not compare?

A5: Thank you for your question! Our method is fundamentally different from GNS and MGN. In GNS, the authors treat positional differences as velocities and concatenate them with the original features as node inputs. The model then predicts accelerations to compute velocities and positions step by step. Additionally, GNS lacks consideration for physical symmetry in 3D space, which can negatively affect the model's generalization ability. In MGN, depending on the experimental setup, some directly predict differences, while others concatenate these differences with original features as node inputs. Their approach to spatial equivariance involves converting coordinate features into relative edge features, which may reduce the model's expressiveness.

In contrast, while our method also incorporates differences, we treat the Temporal Difference Graph (TDG) as a global graph that interacts with all temporal graphs. At each stage, the node features in the TDG are propagated and updated. Finally, we do not predict additional differences; instead, we use the updated node features in the TDG, combined with the last temporal graph, to obtain the predicted coordinates. Moreover, our network is designed to conform to E(3) symmetry (3D translation and rotation symmetry), ensuring that the predicted 3D coordinates are independent of the choice of coordinate system. Furthermore, we specifically design the TDG to support equivariance and prove that it satisfies E(3)-equivariance.

In summary, our method is significantly different from GNS and MGN. To further validate its effectiveness, we conducted experiments using GNS on the MD17 dataset with 8 molecules. The results in Table B demonstrate that our method significantly outperforms GNS, highlighting its superiority.

Tabel B: The performance on MD17.

Q6: 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?

A6: Thank you for your suggestion! We have explained GST-V and GST-F in lines 365–368 in the paper. GST-V and GST-F share the same model architecture, with the only difference being that GST-V processes variable-length inputs, while GST-F processes fixed-length inputs. Although ESTAG can also be extended to handle variable-length inputs, its use of Equivariant Temporal Pooling results in more severe accumulation errors. Additionally, we reported in Appendix C.1 (Table 6) the experimental results of converting all methods capable of handling variable-length inputs into this setting. The results show that for ESTAG and most methods, performance declines significantly across multiple molecules when processing variable-length inputs.

Q7: Why can $X_\delta$ reduce the cumulative error in long-term predictions? This should only be a short-term bias term. What is the theoretical basis?

A7: Thank you for your question, which is highly valuable. We discovered the TDG technique incidentally during our experiments. In our initial attempts, we observed that it achieved general performance improvement across various datasets. We also explored incorporating higher-order terms, such as second-order differences, but the experimental results did not meet expectations. The detailed results are as follows in Table C.

Regarding the theoretical explanation, we hypothesize that this phenomenon can be analogized to first-order numerical methods in solving differential equations. Many AI models draw inspiration from first-order numerical methods. For example, the Euler-Maruyama method, commonly used for solving SDEs in recent Diffusion models, is a first-order method.

Although this explanation is not entirely rigorous, we respectfully request reviewers to focus on the engineering value of our TDG technique. We believe that identifying a simple yet effective engineering solution is a meaningful academic contribution.

Table C: The performance of GST when using second-order differences on MD17.

[1] Hoogeboom E et al. Equivariant Diffusion for Molecule Generation in 3D, ICML 2022.

[2] Xu M et al. Equivariant Graph Neural Operator for Modeling 3D Dynamics, ICML 2024.

[3] Wang Y et al. Enhancing geometric representations for molecules with equivariant vector-scalar interactive message passing, Nature Communications 2024.

[4] Wang T et al. Ab initio characterization of protein molecular dynamics with AI2BMD, Nature 2024.

[5] Thomas N et al. Tensor field networks: Rotation-and translation-equivariant neural networks for 3d point clouds, Arxiv 2018.

[6] Liao Y et al. Equiformer: Equivariant graph attention transformer for 3d atomistic graphs, ICLR 2023.

[7] Vaswani A et al. Attention Is All You Need, NIPS 2017.

[8] Touvron H et al. LLaMA: Open and Efficient Foundation Language Models, Arxiv 2023.

Dear Reviewer V53b,

We extend our sincerest gratitude for your time and valuable feedback in reviewing our paper! We have meticulously responded to all of your questions and concerns, and we would like to know if you are satisfied with our responses and willing to reconsider your evaluation and score for our paper.

As the rebuttal deadline is approaching, we sincerely hope to receive your feedback at your earliest convenience.

Best regards,

The Authors

Dear Reviewer V53b,

We have made our best effort to address all your questions and concerns, and we would like to know if you are satisfied with our responses. As the rebuttal phase is nearing its end, we sincerely hope to receive your feedback, which would greatly help us improve our paper. We would be deeply grateful if you could reconsider your evaluation and consider raising the score!

Best regards,

The Authors