Equivariant Spatio-Temporal Attentive Graph Networks to Simulate Physical Dynamics

Summary

This paper introduces an E(3)-invariant temporal attention scheme, calculated with the help of discrete Fourier transform, within the E(3)-equivariant GNN framework. The overall idea of considering higher-order temporal effects in physics is sound, and the formulation appears to be correct. There are a few typos that do not affect the overall scoring, I would recommend the authors do a full proofreading. There may be missing references and potentially a missing benchmark to compare with. I recommend adding these (in **Weaknesses**). Overall, this work is solid, and I recommend accepting it for Neural IPS 2023.

Strengths

The idea presented in the paper is novel, although not groundbreaking. It fills a gap in the existing framework and the direction is practical and meaningful. There are empirical improvements. The illustrations are very easy to follow.

Weaknesses

There are some typos in the paper, such as the missing year in reference [16]. I recommend thorough proofreading. —Lack of previous SOTA for comparison— [1] Chen, Runfa and Han, Jiaqi and Sun, Fuchun and Huang, Wenbing. "Subequivariant Graph Reinforcement Learning in 3D Environments". Link: https://arxiv.org/abs/2305.18951 —Lack of reference for future improvements— One future direction I have, which has already been used in [1], is combining equivariance with multi-scale (MS) GNN, as most industrial-level applications involve huge graphs. Therefore, the following papers should be cited as future works. Note that [2] also combines equivariance with MS, similar to [1]. However, since there are significant differences in graph type and application needs between this paper and [2], it is not suggested to compare them directly (but they should still be cited). [2] Lino, Mario and Fotiadis, Stathi and Bharath, Anil A and Cantwell, Chris D. “Multi-scale rotation-equivariant graph neural networks for unsteady Eulerian fluid dynamics”. Link: https://pubs.aip.org/aip/pof/article/34/8/087110/2847850 [3] Cao, Yadi, Menglei Chai, Minchen Li, and Chenfanfu Jiang. "Efficient learning of mesh-based physical simulation with bi-stride multi-scale graph neural network.". Link: https://openreview.net/forum?id=2Mbo7IEtZW [4] Meire Fortunato, Tobias Pfaff, Peter Wirnsberger, Alexander Pritzel, Peter Battaglia. “MultiScale MeshGraphNets”. Link: https://arxiv.org/abs/2210.00612

Questions

Please refer to the **Weaknesses** section.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

5: You are absolutely certain about your assessment. You are very familiar with the related work and checked the math/other details carefully.

Soundness

4 excellent

Presentation

3 good

Contribution

3 good

Limitations

Please see the comments regarding "combining equivariance with multi-scale" in the **Weaknesses** section. The attention module for higher-order temporal relationships will significantly increase complexity, which may limit the application in industrial scenarios with huge graphs. The author should either acknowledge this limit, and/or analyze potential remedies for this overhead.

Summary

This paper addresses the Markov limitation of previous methods in simulating physical dynamics by treating it as a spatio-temporal prediction task. The authors propose Equivariant Graph Neural Networks (GNNs) to account for the non-Markovian nature of the systems. Additionally, they design three components to extract spatio-temporal features while preserving equivariance. The experiments conducted on three real datasets demonstrate that the proposed method surpasses previous approaches and validate the effectiveness of the three designed components.

Strengths

1. Identification of Markov limitation: The paper recognizes the Markov limitation present in previous methods and appropriately considers the non-Markovian nature of the systems. This approach is well-founded. 2. Equivariant property preservation: The designed components successfully maintain the equivariant property while extracting spatio-temporal features. The paper also provides theoretical evidence of the proposed Equivariant Spatio-Temporal Attentive Graph (ESTAG) being E(3)-equivariant. 3. Comprehensive experimental validation: The paper includes a substantial number of experiments to substantiate the proposed methods. The results consistently demonstrate the superiority of the proposed approach over alternative methods and confirm the effectiveness of the designed Equivariant Discrete Fourie rTransform (EDFT), Equivariant Spatial Module (ESM), and Equivariant Temporal Module (ETM).

Weaknesses

1. Lack of clarity regarding EDFT: The paper does not provide sufficient explanation of how EDFT improves prediction accuracy. It is important to clarify the underlying mechanisms and intuition behind the proposed component. 2. Need for additional experiments: It would be valuable to conduct further experiments to explore the performance of the proposed method in long-term recurrent forecasting. Providing results in such scenarios would enhance the understanding of the model's capabilities.

Questions

Please refer to weaknesses.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

The authors do not address the limitations. Please refer to weaknesses.

Summary

This paper studies the non-Markovian dynamics that often appear in physical systems and proposes a spatio-temporal E(3) equivariant graph network that moves beyond the simple frame-to-frame prediction task. The authors introduce an equivariant feature extraction method based on Fourier Transform, as well as separable equivariant spatial and temporal modules to process spatio-temporal information. They evaluate the proposed method on different benchmarks and vastly outperform frame-to-frame equivariant methods and non-equivariant spatio-temporal GNNs.

Strengths

The paper studies the non-Markovian dynamics in physical systems, an often overlooked yet very important property. It is well-written and easy to follow. The novelties are clear, and the ablation studies support their usefulness. The quantitative results show a massive performance gain from incorporating equivariance and sequence dynamics.

Weaknesses

The paper claims that "we are the first to use equivariant spatio-temporal graph models for physical dynamics simulation", yet it is missing out on a few related works, namely LoCS [1], and more recently, EqMotion [2]. Both works propose equivariant graph networks and focus on sequence-to-sequence prediction for physical systems. Hence, the authors should do a more thorough literature review and adjust their claims. Many of the neural network modules used in the proposed method are not adequately described in the manuscript, and their significance is not tested with an ablation study. For example, the learnable parameters $w_k$ are only briefly described, and their exact form, as well as their usefulness, are unclear. #### References [1] Kofinas, Miltiadis et al. Roto-translated Local Coordinate Frames for Interacting Dynamical Systems. NeurIPS 2021. [2] Xu, Chenxin et al. EqMotion: Equivariant Multi-agent Motion Prediction with Invariant Interaction Reasoning. CVPR 2023.

Questions

1. Following the weaknesses above, a comparison with other spatio-temporal equivariant graph networks would further enhance the credibility of the proposed method. 2. How important is the use of $w_k$, given that these features are further processed during message passing? 3. Since this work focuses on non-Markovian dynamics, an ablation study on the optimal number of past timesteps would be beneficial and insightful.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

2 fair

Presentation

3 good

Contribution

3 good

Limitations

The authors have adequately addressed the limitations.

Summary

This work proposes a novel architecture for predicting physical dynamics. This architecture first extracts the frequency feature of the input dynamics using a new technique. The frequency feature is then processed by spatial-temporal networks to generate predictions for the future dynamic. Through evaluating the novel architecture on multiple datasets ranging from molecular to macro levels, the authors show that their architecture outperform existing models significantly, especially on the molecular benchmark. The authors also present ablation studies to show which module is the most critical in yielding better performance.

Strengths

The architecture proposed in this work (ESTAG) is novel. It is also shown to significantly outperform earlier works in this work on multiple datasets. The paper is well written and easy to read. The ablation studies conducted in this work also shows that the newly proposed frequency computation technique seems to be the most important module in the model. The result on the molecular dataset (MD17) is especially convincing that ESTAG is better than other models.

Weaknesses

There are some details in the evaluation that are not clear. In the MD-17 dataset, the visualization seems to be really hard to differentiate the ESTAG model and the STEGNN model. But why is the MSE difference so big? Is the observed MSE difference in fact important for whatever downstream task that’s important? Or the STEGNN’s result is already good enough? The other two datasets (Protein and Motion datasets) show that ESTAG is still better than other models, but the gap is much smaller than the MD-17 dataset. It is unclear to me why ESTAG on MD-17 yields such a good result but not on others, can the authors provide an explanation on this? Moreover, on the motion dataset, only one person’s trajectory is used. Why only using one person’s trajectory but not all subjects’? Why visualizations on the other two datasets are not provided?

Questions

See the weakness.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

yes

Summary

This paper aims to simulate the physical dynamics with a spatio-temporal attentive graph network. The major contribution is to integrate the concept of spatio-temporal graph neural network with DFT to capture the data dependencies. The proposed model is evaluated on three datasets regarding the molecular-, protein- and macro-level prediciton.

Strengths

1. The paper is well-written and easy to follow. 2. The proposed model is technically solid. 3. The experiments cover three real-world datasets.

Weaknesses

1. My primary concern is technical novelty. This paper adapts STGNN for physical dynamic simulation without significant contributions or novel designs. The paper claims that it is important to relieve from the Markovian assumption, however, it has been widely explored in the literature on STGNNs. 2. The related work should be further investigated. As far as I know, there are many advanced STGNNs that can achieve much higher accuracy than STGCN (which was proposed in 2018). For more details, please refer to a recent survey [1]. 3. More powerful baselines should be discussed and considered as baselines [1]. This paper only compares ESTAN with STGCN in the line of STGNNs, which is not convincing enough. For example, comparing ESTAG with an existing STGNN published in 2022 is not investigated. Reference: [1] Jin, Guangyin, et al. "Spatio-temporal graph neural networks for predictive learning in urban computing: A survey." arXiv preprint arXiv:2303.14483 (2023).

Questions

Please see the weaknesses.

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

2 fair

Limitations

N/A

Reply

Thanks for your reply. I made a mistake in suggesting [A] to you. It should have been Learning Physical Dynamics with Subequivariant Graph Neural Networks (NeurIPS 22), https://arxiv.org/pdf/2210.06876.pdf Considering the time is not enough for you to add any comparison before the discussion period. Now I can only suggest that the correct version be cited now. If you got accepted in the end, you should try comparing to this one in the final revision. Best,

Official Comment by Reviewer Gpx3

I would like to thank the authors for their rebuttal; they have addressed my questions and concerns. I am still positive towards this paper and I think its contribution is significant. I am keeping my score to 6. The authors should include the related works discussed during the rebuttal in the camera ready version, along with any experiments that compare against them.