Spatiotemporal Learning on Cell-embedded Graphs

Authors: Yuan Mi, Qi Wang, Hao Sun
Primary area: learning on time series and dynamical systems
Keywords: Spatiotemporal Dynamics, Graph Learning, Physics-embeded Learning

Abstract:
Data-driven simulation of physical systems has recently kindled significant attention, where many neural models have been developed. In particular, mesh-based graph neural networks (GNNs) have demonstrated significant potential in predicting spatiotemporal dynamics across arbitrary geometric domains. However, the existing node-edge message passing mechanism in GNNs limits the model's representation learning ability. In this paper, we proposed a cell-embedded GNN model (aka, CeGNN) to learn spatiotemporal dynamics with lifted performance. Specifically, we introduce a learnable cell attribution to the node-edge message passing process, which better captures the spatial dependency of regional features. Such a strategy essentially upgrades the local aggregation scheme from first order (e.g., from edge to node) to a higher order (e.g., from volume to edge and then to node), which takes advantage of volumetric information in message passing. Meanwhile, a novel feature-enhanced block is designed to further improve the performance of CeGNN and alleviate the over-smoothness problem, via treating the latent features as basis functions. The extensive experiments on various PDE systems and one real-world dataset demonstrate that CeGNN achieves superior performance compared with other baseline models, significantly reducing the prediction errors on several PDE systems.

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

Reviews and discussion:
<thread>
creation date:
1730121111227

thread_type:
official_review

<thread_object>
creator: reviewer_kt6u

<content>
summary:
This paper proposes an end-to-end graph-based framework called CeGNN to address limitations of existing Graph Neural Networks in learning complex spatiotemporal dynamics, particularly the over-smoothing issue, and aims to enhance prediction accuracy and generalization ability. The authors introduce two key components: Cell-embedded MPNN block and Feature-Enhanced (FE) block. Through experiments on several PDE systems, the paper demonstrates that CeGNN outperforms other baseline models.

strengths:
- The paper is well-organized and easy to follow.

- The authors present abundant experimental results and visualizations to validate their ideas.

- CeGNN achieves superior performance compared to the compared methods.

weaknesses:
- The baseline methods are relatively weak. The authors did not include recent advancements in the field from 2023-2024, raising concerns about the effectiveness of the proposed method.

- Modeling with higher-order graphs is a widely studied topic. Can the authors more explicitly summarize the contributions of CellMPNN compared to existing approaches?

- The paper lacks a theoretical discussion on the effectiveness of CeGNN. Can the authors discuss the source of CeGNN's effectiveness from a theoretical perspective?

- FE modules are not clearly defined.

questions:
Please address the weaknesses mentioned above.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730147142985

thread_type:
official_review

<thread_object>
creator: reviewer_1r7x

<content>
summary:
The paper presents a new model, the Cell-Embedded Graph Neural Network (CeGNN), for simulating spatiotemporal dynamics across different physical domains. CeGNN introduces learnable cell attributions to the traditional node-edge message-passing process, upgrading it to a higher-order scheme that captures volumetric information and improves spatial dependency learning. Additionally, the Feature-Enhanced (FE) block enriches feature representations, tackling the over-smoothness issue common in Graph Neural Networks (GNNs). Extensive experiments demonstrate that CeGNN achieves superior performance and generalization in predicting physical dynamics, particularly for Partial Differential Equations (PDEs) and real-world datasets.

strengths:
Good empirical results.

weaknesses:
The main weakness is lack of novelty. The main idea of this paper is the proposal of cell-attribution. In the field of topological learning, there have been several prior works proposing the idea of higher-order message passing, cell / simplicial complex neural networks. Please check the following literature:

(1) Topological Deep Learning: Going Beyond Graph Data (this is a great survey of Topological Deep Learning)
https://arxiv.org/abs/2206.00606

(2) Cell Complex Neural Networks
https://arxiv.org/abs/2010.00743

(2) Weisfeiler and Lehman Go Topological: Message Passing Simplicial Networks (for simplical complexes)
https://arxiv.org/abs/2103.03212

However, application of these topological methods in the domain of learning physical systems is new.

questions:
Please address the weakness in novelty that I have raised!
</content>
</thread_object>

</thread>
<thread>
creation date:
1730311571251

thread_type:
official_review

<thread_object>
creator: reviewer_rgev

<content>
summary:
The authors proposed a cell-embedded GNN model (aka,CeGNN) to learn spatio-temporal dynamics. They claim that their learnable cell attribution to the node-edge message passing process better captures the spatial dependency of regional features.

strengths:
1. The paper is easy to follow.
2. We can see detailed description of the technical details
3. This paper touches the core problem in this area.

weaknesses:
1. The most puzzling aspect of this paper is that the discussion of related work and the selection of baselines are all based on studies from before 2022. In fact, there have been many breakthrough studies on both graph-based methods and other neural operator approaches in recent years [1].

2. This work appears more like a straightforward extension of MP-PDE, both in terms of methodology and experiments. The paper proposes a cell-based method for extracting spatial relationships, but how much improvement could be observed if this feature were integrated into MP-PDE?

3. The main experimental results are somewhat confusing. Since the code is not available, it is unclear whether the training data was generated from the authors' own simulations or from public datasets, and what the training dataset size is. If the data is self-generated, the comparison with a few simple baselines is not convincing. Furthermore, the authors mention long-term simulations, yet all experiments are based on one-step predictions, which is clearly insufficient.

4. Regarding the core innovation of this paper, the cell feature is merely a learnable feature initialized by the distance to the cell center. Can its significance be verified by theoretical analysis or by measuring the distance between cell features and node features? The benefit here might simply be from adding position awareness, which makes the model fit specific data better. It could even be considered to replace the distance to the cell center with the distance to the nearest PDE boundary for each point, which might also yield improvements.

[1] Wang, Haixin, et al. "Recent Advances on Machine Learning for Computational Fluid Dynamics: A Survey." arXiv preprint arXiv:2408.12171 (2024).

questions:
See weakness above
</content>
</thread_object>

</thread>
<thread>
creation date:
1730322747198

thread_type:
official_review

<thread_object>
creator: reviewer_s14d

<content>
summary:
In general, this paper tackles interesting and meaningful problems governed by PDE. It is well written and the results are sound.

strengths:
Ablation study is pretty great to justify the proposed framework.

weaknesses:
The methodology of feature-enhanced is not sufficient, the authors should write down in the appendix more equations with more explanation. Most importantly, why do authors propose such Algorithm 1, is there any physical or mathmatical meaning/inspiration? or any hypothesis? It would be better to show the train of thoughts of how did author propose this FE instead of just showing its working better.

questions:
See weaknesses
</content>
</thread_object>

</thread>
<thread>
creation date:
1732153646431

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General reply

comment:
Dear Reviewers:

We would like to thank you for your constructive comments, which are very helpful in improving our paper. We are pleased that the reviewers recognized the novelty and effectiveness of our work. In particular, we thank the reviewers for recognizing the *novelty* (S14d, rgEv, and KT6u) and *effectiveness* (S14d, 1r7X, and KT6u) of our method.

We have addressed all the concerns in each individual rebuttal and summarized as follows. Comprehensive revisions and adjustments (indicated in red color) have also been made in the revised paper (please see the **updated .pdf file**).

- We updated the **FE block in Figure 2** to further explain its process and redefined **Algorithm 1** (placed in Appendix).
- We added some new related works and discussed them in **Section 2**.
- We added some results of new baselines on all benchmarks in **Table 1** and a discussion in **Subsection 4.3**.
- More details of new baselines were added in **Appendix Subsection D.2, Table S3, S10-S12**.
- A primary theoretical deduction of the FE block was updated in **Subsubsection 3.1.1**.
- More theoretical analysises of the FE block were added in **Appendix Subsection C.3**.
- A primary theoretical deduction of the CellMPNN block was updated in **Subsubsection 3.1.2**.
- More theoretical analysises of the CellMPNN block were added in **Appendix Subsection C.1-C.2**.
- More detailed results and discussions were added in **Appendix Subsubsection D.3.3-D.3.5, Appendix Table S14-S17**.

Please do feel free to let us know if you have any further questions.

Thank you very much.

Best regards,

The Authors of the Paper
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper proposes an encode-process-decode framework leveraging graph neural networks (GNNs) to model the spatio-temporal dynamics of partial differential equations (PDEs). It extends the classical message-passing paradigm in GNNs by introducing two key components: (i) leveraging higher-order features between graph nodes to capture spatial dependencies beyond immediate neighbor edges, and (ii) enriching node features through higher-order tensor representations to maintain feature distinction and mitigate the over-smoothing issue caused by successive feature aggregations during message-passing operations. The model is evaluated against baseline models across a variety of datasets and demonstrates improved performance.

The reviewers acknowledge the experimental improvements over the baselines. However, they raised several concerns about the initial version of the manuscript. These include the incremental nature of the technical contribution—several reviewers pointed out the existence of alternative methods for capturing higher-order features in GNNs—and the choice of baselines for comparison.
In response, the authors provided extensive new comparisons during the rebuttal phase, including implementing the proposed components on other backbones, which also demonstrated improved performance. While this significantly improved upon the first version of the paper, the majority of the reviewers maintained their initial concerns, particularly regarding the incremental value of the contribution.