SEGNO: Generalizing Equivariant Graph Neural Networks with Physical Inductive Biases

Thank you for the provided comments and helpful suggestions! We have enhanced our paper accordingly.

Comment a: Lack of Visual Illustration.

Thanks for the instructive suggestion. We have included visualizations of N-body systems in the revised version.

Comment b: Inconsistent color in Figure 4.

Thanks for the comment. We have corrected the colors of Figure 4.

Comment c: Figure 2 description.

Thanks for raising this issue. Figure 2 is a case study of charged N-body systems where EGNN and SEGNO are trained to predict positions after 1000ts. We have added it to the Figure 2 caption in the revised version.

Comment d: Inconsistent Sectioning for Human Motion Experiment.

We apologize for the confusion. We adjust section arrangements to enhance readability in the revised version.

Comment e: Future Direction on Scalability.

Thank you for providing insightful future direction and valuable references! We agree that it is a very interesting point to explore our framework to solve the dynamics of large graphs. Analogous to our work, a feasible solution is to parameterize the target partial differential equation (PDE) via a GNN and incorporate solvers to train the model. We have added the references and will explore them in future work.

Thanks very much for your time and valuable suggestions. In the rebuttal period, we have added additional visualizations of N-body systems in Appendix, corrected Figure 4 colors, Figure 2 description, and inconsistent sectioning for the human motion experiment. We also appreciate the insightful future direction and valuable references. We have added it as future work in the revised version.

Would you mind checking our responses and confirming whether you have any further questions?

Any comments and discussions are welcome!

Thanks for your attention and best regards.

We really appreciate the reviewer for listing the detailed strengths, which have clearly recognized the novelty and contributions of our paper. As requested by the reviewer, we have experimentally shown the running time of all models.

Weakness 1: Running time of models.

Thanks for your suggestion! We evaluate the running time of each model on N-body systems with Telsa T4 GPU and report the average forward time in seconds for 100 samples. The results are listed below:

We can observe that SEGNO's forward time (0.0227s) remains competitive compared to the best baseline SEGNN (0.0315s), indicating its efficiency.

Question 1: Can SEGNO be included in a Markov Chain Monte Carlo (MCMC) framework？

The reviewer has raised an interesting question that would serve as a possible future direction. Although both SEGNO and the MCMC framework (Metropolis-Hastings) aim to explore molecular dynamics, they are different in the following aspects:

• SEGNO primarily focuses on simulating the time-dependent evolution inherent in Newtonian dynamics.
• MCMC framework is targeted towards thermodynamic analysis of molecular systems at equilibrium, rather than dynamic simulations of how systems evolve over time.

Given their different operational domains and objectives in molecular dynamics, integrating SEGNO with the MCMC framework remains a challenging task. However, a possible solution is to replace the target ODE of SEGNO with a Stochastic Differential Equation (SDE). Thus, MCMC framework can be employed to estimate parameters within the SDE. We will explore it in the future.

Question 2: Do the authors envision the potential for SEGNO to be employed in a similarly transferable manner?

Yes! We agree that it is a very interesting and promising point to explore. To investigate the model transferability across different systems, we train models on 5-body Gravity systems and test them on 10-body and 20-body systems. The goal is to predict the positions after 1000ts. The results are as follows:

We can observe that although SEGNO does not train on trajectories of 10-body and 20-body systems, it achieves significantly better performance, demonstrating its transferability across different systems.

Moreover, for transferability across different tasks, since the goal of SEGNO is to model second-order dynamical systems (i.e., Eq.1), it can be naturally generalized to predict other dynamic states such as energy and momentum.

Thanks for your constructive comments! We try our best to address your concerns about the novelty and theoretical findings. As suggested by the reviewer, extensive additional experiments including additional baselines, metrics, and backbones, are conducted to further support our contributions.

On the novelty.

Weakness 1: The idea of combining GNODE with an equivariant GNN is fairly straightforward.

Thanks for the comment. We would like to clarify that our contributions do NOT lie in the combinations of GNODE and equivariant GNNs. Existing studies[1,2] mainly employ a GNN that iteratively updates the latent representation to forecast the average acceleration. We distinguish our novelty and significance in the following aspects:

• Our goal of providing theoretical insights to build Graph Neural ODE is novel and has never been explored before as far as we know. We show that a second-order Graph Neural ODE that iteratively updates positions instead of latent representations, can obtain bounded error of instantaneous acceleration and position through minimizing position discrepancy.
• We thank the reviewer for raising the comparison with the works[1,2]. We validate our findings in novel applications (e.g., molecular and human motion dynamics) where the geometric property is crucial and propose incorporating equivariant GNNs to fulfill this inductive bias. Empirical results demonstrate our model has better generalization ability.
• Thanks for your reference[1]. We agree that it is feasible to change NODE[1] backbones to equivariant GNNs. Nevertheless, it is still not clear whether such replacement affects the equivariant property of models. Our Proposition 3.1 provides a guarantee to build an equivariant GNODE.

Thanks for your valuable references. We have added discussions of these papers in the revised version. We believe such a discussion can clarify the main contribution of our paper.

[1] Unravelling the performance of physics-informed graph neural networks for dynamical systems. NeurIPS, 2022.

[2] Enhancing the inductive biases of graph neural ode for modeling dynamical systems. ICLR, 2023.

[3] Deconstructing the inductive biases of hamiltonian neural networks. ICLR, 2022.

Compared with advanced force field prediction models, such as NequIP and Allegro.

Thanks for your suggestions. We would like to clarify that the goal of baselines such as DimeNet, NequIP, and Allegro, is force field prediction, which is different from our task as follows:

• Force field prediction is a static task that mainly takes atom positions as input, and outputs molecular energy and atom forces. To ensure the conservation of energy, the force predictions are partial gradients of predicted energy.
• Trajectory prediction is a dynamic task that mainly takes atom positions and velocity as input, and outputs future atom positions.

Weakness 3: It is not clear why advanced baselines, such as DimeNet, NequIP, Allegro, MACE, BOTNet, Equiformer, etc., are not used.

Thanks for your suggestion. We add empirical comparison on MD22 with official implementation Nequip[1] and Allegro[2]. The models take atom positions and numbers as input and optimize the errors between geometric outputs (generally treated as predicted forces) and true positions. We try our best to tune the model and the results are as follows:

It can be observed they do not perform well in this setting due to the lack of dynamic information.

Question 5: Backbones such as NequIP would be better for such cases.

We are grateful for advising us to consider NequIP as backbones, which is an important direction to improve model performance. However, as shown in the above results, it is non-trivial to incorporate dynamic geometric features (e.g., velocity) in the framework of irreducible representations. The ways that map velocity into high-order irreducible representations, design interaction modules, and decode dynamic states from irreducible representations are still an open problem. Thus, we are interested in exploring them in future works. We provide the results of the EGNN backbone in the following responses.

On the theoretical findings.

In the following, we provide additional details to eliminate the confusion. It's important to note that these details do Not affect the correctness or contributions of our theoretical findings.

Question 2: Why is Theorem 4.1 a property of SEGNO?

Apologize for the confusion. We agree that Theorem 4.1 is a general statement regarding the uniqueness of solutions for dynamic systems. This statement serves as the prerequisite for Proposition 4.2 that SEGNO is capable of approximating the instantaneous acceleration of such a unique trajectory. We acknowledge that the current section title may potentially lead to misunderstandings. Thus we change the title as follows:

• Section 4 -> SEGNO Analysis
• Theorem 4.1 -> Lemma 4.1

Question 1: Authors may rephrase the text accurately approximate in the proof
Question 3: Why Proposition 4.2 warrants a separate proof?

We apologize for the confusion caused by the omission of some technical details in the proof. The non-trivial aspects of the proof mainly lie in demonstrating that minimizing the position loss $|| \mathbf{q}\_{\theta}^{(t\_{1})} - \mathbf{q}^{(t\_{1})} ||\_p$ can enable neural networks to approximate the acceleration $f(\mathbf{q}^{(t)})$. Previous studies either directly predict the tartget position $\mathbf{q}^{(t\_{1})}$ by minimizing position loss $|| \mathbf{q}\_{\theta}^{(t\_{1})} - \mathbf{q}^{(t\_{1})} ||\_{p}$[1,2] or approximate acceleration $f(\mathbf{q}^{(t)}, \mathbf{h})$ using acceleration loss $|| f\_{\theta}(\mathbf{q}\_{\theta}^{(t)}) - f(\mathbf{q}^{(t)})||\_{p}$[3].

In contrast, SEGNO has the ability to convert position loss into acceleration loss. In the meantime, it's a common practice to leverage the universal approximation theorem to establish the universal approximation capability of NNs within deep learning community[4,5,6]. Therefore, according to the aforementioned analysis and universal approximation theorem, it can be shown that SEGNO is capable of approximating acceleration $f(\mathbf{q}^{(t)}, \mathbf{h})$ with position loss.

Due to the space constraints, we provide only a concise discussion here. For the complete mathematical derivation, we kindly refer the reviewer to the proof presented in the revised version. Additionally, we would like to re-emphasize that, while Proposition 4.2 seems intuitive, this property has made significant contributions in the following aspects:

• To the best of our knowledge, it's the first theoretical justification that NNs equipped with SEGNO framework possess the capability to approximate the unique acceleration using position loss.
• Empirically, it enables SEGNO to recover latent trajectories between input and output states, as illustrated in Fig. 2.

Regarding your concern about the term accurately approximate, we agree that it introduces ambiguity pertaining to the level of accuracy. To address the vagueness, we'll rephrase the text to replace accurately approximate with simple approximate.

[1] Vıctor Garcia Satorras, Emiel Hoogeboom, and Max Welling. E (n) equivariant graph neural networks. ICML, 2021.

[2] Wenbing Huang, Jiaqi Han, Yu Rong, Tingyang Xu, Fuchun Sun, and Junzhou Huang. Equivariant graph mechanics networks with constraints. ICLR, 2022.

[3] Alvaro Sanchez-Gonzalez, Jonathan Godwin, Tobias Pfaff, Rex Ying, Jure Leskovec, and Peter Battaglia. Learning to simulate complex physics with graph networks. ICML, 2020.

[4] Universal Invariant and Equivariant Graph Neural Networks. NeurIPS, 2019.

[5] A Universal Approximation Theorem of Deep Neural Networks for Expressing Probability Distributions. NeurIPS, 2020.

[6] Deep network approximation: achieving arbitrary accuracy with fixed number of neurons. JMLR, 2022.

Additional backbone.

Weakness 2: Evaluation on additional backbone architectures is missing.

Thanks for the suggestion. We report the results of EGNN and GMN backbones as follows:

We can observe that employing the GMN backbone achieves better results, which could be attributed to its multi-channel feature representation.

Additional metrics and physics-informed GNN baselines.

Weakness 4: There are several metrics such as energy violation error, momentum error, rollout error etc. in the literature.
Weakness 5 & Question 4: Comparison with physics-informed GNNs.

Thank you for raising the comparison with physics-informed GNNs. As suggested by the reviewer, we thoroughly compare the rollout position, energy, and momentum errors of our SEGNO with those of GNODE and HGNN on N-body systems [1]. GNODE and HGNN aim to predict future acceleration and Hamiltonian respectively. All models are trained by minimizing the errors between the prediction and dynamic states (e.g., positions or acceleration) after 100 timesteps (a 0.001 second as 1 timestep). The results are shown as follows:

On Charged N-body systems:

On Gravitational N-body systems:

We can observe that SEGNO consistently outperforms GNODE and HGNN in rollout errors. Although SEGNO is only trained on the system positions, it achieves comparable performance on energy and momentum errors with GNODE, which further validates our findings.

[1] Unravelling the performance of physics-informed graph neural networks for dynamical systems. NeurIPS, 2022.

Thanks very much for your time and valuable comments. As the discussion will close in a few hours, we would be grateful if you could allocate some time to review our response.

We understand that you have a multitude of responsibilities, thus we have summarized our response as follows:

• We add discussions of existing GNODEs to clarify the main contribution of our paper (Response-(1/3)).
• We add advanced force field prediction models as baselines (Response-(1/3)).
• We empirically compare our model with physics-informed GNN baselines w.r.t. the rollout position, energy, and momentum errors (Response-(2/3)).
• We evaluate the performance of different backbones (Response-(2/3)).
• We provide additional details to eliminate the confusion about our theoretical findings (Response-(3/3)).

Would you mind checking our responses and confirming whether you have any further questions?

Any comments and discussions are welcome!

Thanks for your attention and best regards.

Dear Area Chair,

Thanks for your kind reminder.

Best regards,

Authors