Equivariant Neural Simulators for Stochastic Spatiotemporal Dynamics

Summary

This paper proposed to incorporate symmetries into probabilistic models to simulate stochastic dynamics over trajectories. In particular, the authors developed EPNS, an equivariant version of sequential CDVAE for the generation of the state transition. The results on two tasks including stochastic n-body system and stochastic cellular dynamics, show that EPNS considerably outperforms existing neural network-based methods for probabilistic simulation in terms of simulation quality, data efficiency, rollout stability, and uncertainty quantification.

Strengths

1. It is valuable to design equivariant probabilistic neural simulators against existing deterministic counterparts. It is true that random effects exist in common systems and should be characterized. 2. Lemma 1 provides convincing derivations for the formulation of equivariant CDVAE. It is interesting to see that there are two different ways to derive equivariant models: invariant latent and equivariant latent. Besides, severe deviation over long rollout simulations has always been a tricky problem. To mitigate this phenomenon, the paper adopts a heuristic method by using the ground truth of the next state to steer the sampling progress， which is inspiring. 3. This paper is well-organized, first introducing the general formulation and then detailing the specific architectures for two interesting applications. I enjoy the writing. The experimental evaluations, although insufficient, is appropriately demonstrated to support certain claims of the authors.

Weaknesses

1. My biggest concern is that one important reference [A] is not cited and discussed. The authors claim that “current probabilistic methods do not incorporate general domain symmetries into their model architectures”, which is incorrect if checking [A]. The work by [A] has developed equivariant Gaussian Process (GP) to model stochastic fields, including scalar, vector and stacked fields. Equivariant GP is applicable to the task of modeling system evolutions discussed in this paper. We can directly consider p(x_{t+1}|x_t) as the prediction of vector fields studied in [A]. More importantly, when referring to [A], many claims of the authors in Introduction (Lines 47-51, Lines 55-57) and related works (Line 95-97) are no longer correct, and [A] should be at least considered as an indispensable baseline in the experiments. It is unclear how much novelty remains and how many extra efforts should be made to reflect the modification of discussing [A]. 2. For the experimental evaluations, several crucial contribution points are NOT justified sufficiently. First of all, the equivariant deterministic methods should be tested, by which we can observe why we need stochastic modeling. Second, besides the reference [A] mentioned above, the authors are suggested to compare against typical probabilistic models (such as ODEVAE) with trivial equivariant modification, which is able to better support the benefit of the proposed designed model. Third, in comparison with using the latent model z, one can directly model the transition probability p(x_{t+1}|x_t) straightforwardly via multiple equivariant layers. How does this baseline perform? 3. The authors proposed to use invariant latent for Celestial dynamics, and equivariant latent for cellular dynamics. Certain ablations are suggested to provide to support this design. 4. The conclusion of Lemma 1 is interesting. It provides sufficient conditions to maintain the symmetries, but are these conditions also necessary? 5. Would you please explain more on why EGNN is unstable for the Celestial dynamics task? 6. Is it possible to combine invariant latent and equivariant latent into one single model. In this case, does it mean that invariant latent corresponds to invariant features in EGNN and equivariant latent corresponds to equivariant features in EGNN? 7. How to generate stochastic ground truth trajectories in the experiments? How to compute quantile interval of the observations in Figure 5. [A] Equivariant Learning of Stochastic Fields: Gaussian Processes and Steerable Conditional Neural Processes, ICML 2021.

Questions

See the weakness part above.

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

The authors have adequately discussed the limitations.

Summary

This paper introduces a general framework named Equivariant Probabilistic Neural Simulation (EPNS) for equivariant stochastic simulation with neural simulators. The framework incorporates domain equivariances and utilizes an autoregressive probabilistic approach. The authors prove that, under equivariance assumptions on the neural network layers, the probability distribution over trajectories remains equivariant. The proposed framework is applied to two scenarios: a stochastic n-body system with celestial dynamics and stochastic cellular dynamics on a grid domain. Additionally, the authors propose a novel message passing graph neural architecture that maintains equivariance to permutations of cell indices, motivated by stochastic cellular dynamics problem. Experiments showcase multiple facets of the benefits of this framework.

Strengths

S1. The paper presents a method for learning stochastic equivariant distributions over trajectories, which constitutes a valuable contribution to neural simulation. Additionally, the authors introduce a novel message passing graph neural network that enables equivariant embeddings in domains containing lattices and sets. S2. The experiments are comprehensive and demonstrate several benefits of the framework, including improved data efficiency, uncertainty quantification, and rollout stability. S3. The paper's presentation is excellent, with a detailed appendix and code that facilitates reproduction of results. The results are comprehensive, and the insights are well-described.

Weaknesses

W1. The authors state the benefits of learning equivariant distributions over equivariant function approximations. However, it is unclear whether any of their experiments support this claim beyond enabling uncertainty quantification. W2. The authors did not explicitly state whether the simulations used for the experiments are toy examples or scenarios that are of significant interest to domain experts. It would be beneficial for the paper to provide clarification on this aspect. Specifying the nature and relevance of the simulations would help us assess the potential impact of the developed methods on scientific applications. This information is crucial for understanding the current progress in developing methods that have practical implications in the scientific community.

Questions

Q1. While I understand how equivariance can help in maintaining cell clusters and representing the true number of clusters, I am unclear about its connection to kinetic energy in celestial dynamics. Could the authors provide further insight into how equivariance relates to kinetic energy in this context? Q2. The authors did not cite any papers when selecting the metrics presented. Could the authors explain the rationale behind not basing their metric selection on previous work in neural simulation? Q3. How can EPNS be extended to handle non-Markovian dynamics? What would be the potential benefits of such an extension? Q4. Please comment to W2.

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

4 excellent

Presentation

4 excellent

Contribution

3 good

Limitations

Limitations were mentioned. See W2 for another potential limitation.

Summary

This work presents a new framework for autoregressive probabilistic modeling called Equivariant Probabilistic Neural Simulation (EPNS) that produces equivariant distributions over system evolutions. The authors demonstrate that this method outperforms existing neural network-based methods for probabilistic simulation, improving simulation quality, data efficiency, rollout stability, and uncertainty quantification. The paper also discusses the benefits of incorporating domain symmetries in deterministic neural simulators and provides specific examples of domains where EPNS could be applied for efficient and effective data-driven probabilistic simulation.

Strengths

Authors attempt to solve a challenging problem of simulating systems following a probabilistic simulation. They present an autoregressive approach toward probabilistic simulations combining equivariance. This can be used for several problems such as Langevin dynamics, Brownian motion to name a few.

Weaknesses

There are several weaknesses for the work as outlined below. 1. **Simple systems:** Authors have considered very simple toy systems such as 3-body gravitational system, and a simplified cellular dynamics, a lattice based system. There are several examples of realistic probabilistic simulations such as Langevin or Brownian dynamics; these are not considered. This limits the application to toy systems and validity of such approach for larger realistic systems are questionable. 2. **Generalizability:** The work is demonstrated on a 3-body gravitational system. The GNN nature of the approach allows natural zero-generalizability to larger system sizes. These are not explored making the application of the approach to real-world problems limited. Similarly, due to the data-driven approach, it is not clear whether the approach can be generalized to unseen noise conditions for the same system. That is, for the gravitational system, if the interactions with the dust particles increased or for the cellular the systems, if the temperature increased, will the approach be applicable is not clear. 3. **Metrics:** This approach is presented as a data-driven approach to learn the dynamics. However, the ability of the model to preserve the physical quantities meaningfully is not explored. For instance, the forces on each particle in the rollout, how do they compare with those in the ground truth in a statistical fashion? Similarly, the phase space explored by the rollout, do they remain same ensuring statistical similarity of the trajectory. Indeed, the authors have shown the $LL$ and $D_{KS}$, but more elaborate error metrics can provide insights into the nature of learned simulations. 4. **Baselines:** It is not clear why a graph architecture is chosen especially when there is no generalizability to larger system sizes explored. To this extent, an additional MLP-based baseline can give insights into whether a GNN architecture provides superior results or not.

Questions

There are several concerns regarding the work as outlined below. 1. There are several equivariant architectures including the EGNN authors have mentioned. The choice of the specific GNN which is not truly state-of-the-art for dynamics simulation is not clear. Authors mentioned that EGNN is not stable while training. EGNN has one of the simplest architectures among the equivariant GNNs and it is not clear why such an architecture would be unstable. Authors should clarify this. 2. The reason for the choice of a graph for 3-body system is not clear. Why not an MLP for such a simple system? 3. Can the graph architecture be generalized to larger systems, say 30 or 300 particles? It is suggested to demonstrate this ability. Otherwise, the approach remains applicable only to toy systems. 4. What are the characteristics of the learned dynamics? Can the drift and diffusion term of the SDE be decoupled? If yes, how well is the variance of the drift term captured by the simulator? If not, how can the ability of the model to realistically capture the simulation be established. 5. Authors discuss about the stability of the simulation. Indeed, autoregressive models are known to be unstable. Can the authors discuss about the stability of the learned simulation with respect to the length of the training trajectory or inference trajectory? That is, can they say that the inference trajectory will remain stable for the same length as the training trajectory or lower or more $x$ times more. This would be helpful to establish the usefulness of the ML-model 6. What about the data-efficiency of the approach? Does it continue to become better with increasing data. If yes, by how much? Have authors investigated the data efficiency of the model? 7. Finally, the model is not interpretable and not generalizable to larger system sizes. For instance, the model cannot give any insights about the physics or dynamics of the system in terms of the drift or the diffusion term. This means it can simulate the same system on which it is trained. This would be useful only if the model is computationally much more economical than the ground truth. Is that the case? Authors have mentioned that the GNN is computationally expensive. In that case, where do the authors think the approach would be useful practically?

Rating

4: Borderline reject: Technically solid paper where reasons to reject, e.g., limited evaluation, outweigh reasons to accept, e.g., good 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

3 good

Limitations

While some of the limitations are briefly mentioned, there are a large number of limitations for the present work as discussed in the weaknesses and questions section. Authors should either address those or include those in the manuscript as limitations.

Summary

The paper presents EPNS, a framework for equivariance in autoregressive probabilistic multidimensional dynamics modeling. The framework is general such that one can use existing equivariant NN architectures; for the task of cellular dynamics forecasting the authors also introduce a ConvNet GNN. The models are shown to outperform current SOTA in both accuracy and data efficiency.

Strengths

The framework and algorithm proposed is sound and general, which I can see being widely used in the spatiotemporal forecasting community. The authors did a great job outlining and explaining the EPNS algorithm step by step. The illustrations are great and helpful as well. The method even addressed the long rollout deviation problem common to autoregressive models. Overall, the algorithm is convincingly effective and practical.

Weaknesses

- It was a little misleading when the authors say that they theoretically proved that EPNS produces equivariant distributions over trajectories. From my understanding, the equivariance comes from the paper's formulation by construction. Maybe a better way to word the contribution is something like "Our framework is able to construct equivariant distributions over the entire trajectory autoregressively from equivariant single-step predictions"? - There is no substantial contributions on model architecture. I personally found the framework and training setup to be insightful enough to make an impact on the forecasting community, but one could argue it's just a new way to use existing models on new datasets.

Questions

1. The paper mentions no work has considered equivariant probabilistic modeling of spatiotemporal dynamics. I can think of the following papers that tackles similar settings as in this paper: - Köhler, Jonas, Leon Klein, and Frank Noé. "Equivariant flows: sampling configurations for multi-body systems with symmetric energies." arXiv preprint arXiv:1910.00753 (2019). - Garcia Satorras, Victor, Emiel Hoogeboom, Fabian Fuchs, Ingmar Posner, and Max Welling. "E (n) equivariant normalizing flows." Advances in Neural Information Processing Systems 34 (2021): 4181-4192. - Sun, Sophia Huiwen, Robin Walters, Jinxi Li, and Rose Yu. "Probabilistic Symmetry for Multi-Agent Dynamics." In Learning for Dynamics and Control Conference, pp. 1231-1244. PMLR, 2023. I believe your framework is more general w.r.t. symmetry groups, but it would be good to cite these works as well to provide context. 2. How are the p-values in table 2 calculated, and how should we interpret them? 3. Some minor questions on writing: - Is it a typo on line 143? equivariant with regards to $h^t$? - Line 277: "meaning that new simulations also match better to the ground truth in terms of the examined properties" <- what does 'new simulation' and 'examined properties' refer to here?

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 author provides a detailed discussion of limitations in the conclusion section of the paper.

Thank you

Thank your for your detail responsese. W1: I am satisfied with the answer. I hope the authors should consider the comparison with the mentioned reference [A], as equivariant stochastic models have been explored before. W2 and W3: I notice the new results on equivariant deterministic methods and the equivariant or invariant latent. That is great. However, I would expect more explanations on why stochastic methods are beneficial. The authors should explain more on this point since it is indeed the main motivation and contribution of this paper. As for the equivariant baselines, I understand that implementing equivariant ODEVAE could be difficult. But the authors should do more on taking other trivial equivariant probabilistic models into comparison. W4: Thank you. W5: The explanations should be included in the revised paper. W6: Thank you. W7: More details should be included in the main paper. Overall, the authors’ responses are helpful and positive. I am willing to increase my score. Given the technical novelty and contributions are not that significant, 5 is the best I can give.

Thank you

Thank you for the detailed response. I would suggest authors modify the manuscript's limitations and future work aspects to reflect these discussions. I have raised my score. While I agree that the 20-body system is reasonable, the 80x80 grid doesn't truly represent a challenging high-dimensional problem. This is because a practical situation will not have a fixed 80x80 grid, and this volume will change. A problem of 6400 particles moving in space would be much more challenging than the 80x80 grid, although the dimensionality of each particle is only 3 and not 64. A realistic application of such stochastic simulators would require dealing much more particles (of the order of 100s or thousands) with significantly more complex interactions. Authors may include this in the limitations section if they agree.