Spatiotemporal Learning on Cell-embedded Graphs

Thanks for your constructive comments!

Weakness 1: Add recent works as new baselines.

Reply: Great suggestion! We have added results of comparison with several other baseline models in the revised paper (see Subsection 4.3 on Page 7, Table 1 on Page 9, and Appendix Tables S3, S10-S12 on Pages 19-21), namely,

• The additional experimental results are added in Table 1.
• A discussion about the baseline comparision results is added in Subsection 4.3.
• Details of the new baselines, e.g., the hyperparameters, have been added in Appendix Subsection D.2, Appendix Tables S3, S10-S12.

Specifically, we have considered three additional baselines (aka, FFNO, Geo-FNO, Transolver) to compare with the performance of CeGNN on our all benchmarks. A summary of these models is shown in the following Table F. The results show that Geo-FNO [1] and Transolver [3] perform poorly on the limited training dataset, which is significantly smaller, by at least an order of magnitude less, than the hundreds or thousands of datasets typically used for such models. We have added the new results to our revised paper, referring to Table 1 (Page 9).

Table F: RMSE metrics for CeGNN and other models.

References:

[1] Li, et al. A. Fourier neural operator with learned deformations for pdes on general geometries. arXiv:2207.05209, 2022.

[2] Tran, et al. Factorized fourier neural operators. In ICLR, 2023.

[3] Wu, et al. Transolver: A fast transformer solver for pdes on general geometries. In ICML, 2024.

Weakness 2a: Summarize the contributions of CellMPNN block from a theoretical perspective.

Reply: Great comment! Generally, the traditional message passing (MP) mechanism can be regarded as a refinement on a discrete space, analogous to an interpolation operation, which implies that edges are essentially interpolated from nodes. A MP mechanism introducing the cell further enhances the refinement of the discrete space (secondary refinement), thereby reducing the magnitude of discretization errors spatially, paving the way for its application in complex graph structures. Please see below Definition 1 on cell in graph and Corollary 1 on expressive power of cell-embedded MP. Therefore, we proposed a new two-level cell-embedded mechanism to process the message on graphs, which forms the key novelty of the proposed model.

Definition 1 (Cell in Graph): Let $G = (V, E)$ be a graph, where $V$ is the set of nodes $\mathbf{v}$ and $E \subseteq V \times V$ is the set of edges. A cell in $G$ is a subset of nodes $C \subseteq V$, such that the nodes in $C$ form a complete subgraph (clique) or satisfy predefined structural relationships. In particular, a $k$-cell $C_k$ in a graph $G$ contains $k+1$ nodes, where $\forall i, j \in C_k$, $(\mathbf{v}_i, \mathbf{v}_j) \in E$, representing various structures, such as node ($k=0$), edge ($k=1$), triangle ($k=2$), tetrahedron ($k=3$), and so on.

Corollary 1 (Expressive Power): Given a graph $G$ including many $k$-cell ($k=0,1,2,\dots$), there exists a cell-based scheme that is more expressive than Weisfeiler-Lehman (WL) tests in distinguishing non-isomorphic graphs (see the proof in Appendix Subsection C.2 of the revised paper (Pages 16-17)).

Weakness 2b: Summarize the contributions of CellMPNN block from an experimental verification perspective.

Reply: In particular, a summary of CellMPNN block on the performance of some graph-based baselines across all benchmarks is shown in the following Table H. The results in Table H show the positive effect of CellMPNN block. We have added the new results to our revised paper, referring to Appendix Table S14 in the revised paper (Page 22).

Table H: Summary of CellMPNN block on the performance of graph-based baselines across all benchmarks

In our study, we also attempted to enhance feature distinguishability through a three-level update structure ($cell \rightarrow edge \rightarrow node$). Since there are only node labels for the supervised learning, this three-level message passing mechanism failed to yield performance lift which might be due to the introduction of redundant path of information passing. Therefore, we adapt two-level message passing sequence to reduce the high coupling degree and avoid the limitation for additional predefined special complexes. A comparison result between two-level and three-level mechanisms is shown in the following Table G. The results in Table G demonstrate that the three-level mechanism underperforms compared with the proposed two-level mechanism in all cases. We have added the new results to our revised paper. Please refer to Appendix Table S17 in the revised paper (Page 23).

Table G: Comparison result between two-level and three-level mechanisms

Weakness 4: Discuss the CeGNN's effectiveness from a theoretical perspective.

Reply: Great suggestion! In fact, the novelty of CeGNN includes the two-level cell-embedded message passing mechanism and the unique feature-enhanced (FE) block. Given the cell-embedded mechanism that better captures the spatial dependency and the FE block that further learns the distinguishable features, CeGNN achieves superior performance compared with other baseline models across all benchmark datasets, as shown in Table 1 in the revised paper (Page 9). Moreover, as shown in Table 3 in the revised paper (Page 9), we have provided an ablation study on all benchmarks to assess the contributions of the FE and CellMPNN blocks in CeGNN.

Following your suggestion, we have added additional discussions in the revised paper (see Subsection 3.1 on Pages 4-6, Appendix Section C on Pages 16-17), namely,

• A primary theoretical deduction of FE block was added in Subsubsection 3.1.1.
• A primary theoretical deduction of CellMPNN block was added in Subsubsection 3.1.2.
• Theoretical preliminaries have also been added in Appendix Section C.

In summary, we proposed a new two-level cell-embedded mechanism to process the message on graphs. Please also see our reply to Weakness 2a and 2b about the theoretical and experimental verification perspective of the effectiveness. Moreover, the FE block is to capture nonlinear dependencies, enhancing the model's representation power. More detailed theoretical discussion about FE block can be found in our reply to Weakness 3.

Overall, we appreciate your constructive comments and suggestions. Please let us know if you have any other questions. We look forward to your feedback!

Weakness 3: Unclearly defined FE modules.

Reply: Great comment! Following your suggestion, we have added detailed explanation in the revised paper (see Section 3.1.1 on Pages 4-5 and Appendix Section C.3 on Page 17).

• We update the FE block in Figure 2 to further explain its process and redefine Algorithm 1 (placed in Appendix).
• A primary theoretical deduction is updated in Subsubsection 3.1.1.
• More theoretical knowledge have also been added in Appendix Section C.3.

Concretely, the FE block is inspired by interaction models in physics and mathematically hypothesized to capture nonlinear dependencies, enhancing the model's representation power.

Hypothesis. The FE block's hypothesis could be framed as: By capturing second-order interactions between latent features and applying selective filtering, the model can better represent complex structures or relationships in the data.

Physical Analogy. In physics, the outer product and second-order terms are often used to model interactions, such as stress tensors in mechanics or pairwise correlations in quantum mechanics. Here, the module could draw an analogy to systems where interactions between individual components (features) are crucial to the overall behavior.

Process Overview. In detail, it regards the node latent feature $\overline{\mathbf{h}}_ i \in\mathbb{R}^{D\times 1}$ as basis and builds a higher-order tensor feature $\mathbf{H}_ {i} \in\mathbb{R}^{D\times D}$ via an outer-product operation, e.g., $\overline{\mathbf{h}}_ {i}\otimes\overline{\mathbf{h}}_ {i}$. This process creates abundant second-order nonlinear terms to enrich the feature map. We then use a mask operation with $\mathbf{M}\in\mathbb{R}^{D\times D}$ to randomly sample these terms, filtering the appropriate information by a learnable weight tensor $\mathbf{W}\in\mathbb{R}^{D\times D\times D}$ to enhance the model's representation capacity.

Outer Product as Basis Expansion. The outer product operation $\otimes$ on the reshaped feature map $\overline{\mathbf{h}}_ {i}\in\mathbb{R}^{D\times 1}$ expands the original latent feature space into a higher-order tensor space. This expansion introduces second-order terms (e.g., $\alpha\beta$ for $\alpha, \beta \in \overline{\mathbf{h}}_{i}$), which can capture interactions between individual components of the original feature $\mathbf{h} _{i} \in\mathbb{R}^{D}$. Mathematically, the second-order tensor reads $\overline{\mathbf{h}} _{i} \otimes\overline{\mathbf{h}} _{i}$.This operation creates a richer feature map with cross-term interactions that may not be explicitly encoded in the original latent space.

Lemma 1 (Nonlinear Representation): The second-order terms $\alpha\beta$ can model nonlinear dependencies between features. This is particularly useful for capturing complex interactions that linear transformations (e.g., simple dot products) might overlook.

Definition 2: The FE block expands the latent feature $\overline{\mathbf{h}}_ i \in\mathbb{R}^{D\times 1}$ of node $i$ into a higher-order tensor space using an outer product: $\mathbf{H}_ {i}=\overline{\mathbf{h}}_ {i} \otimes\overline{\mathbf{h}}_ {i}$, where $\mathbf{H}_ {i}\in\mathbb{R}^{D\times D}$ is a higher-order feature map.

Regularization via Masking. Masking introduces sparsity in $\mathbf{H}_ {i}$, reducing overfitting. If $\mathbf{M}_ {jk}$ is selected, $\mathbf{M}_ {jk} =1$. Otherwise, $\mathbf{M}_ {jk} =0$. Here $j$ and $k$ index the $M \in\mathbb{R}^{D\times D}$ components.

Theoretically, masking operation serves two purposes: (1) Reducing computational complexity by randomly sampling terms from the higher-order feature space. (2) Regularizing the model by introducing sparsity, which can prevent overfitting in high-dimensional spaces.

Learnable Filtering. The learnable weight tensor $\mathbf{W}\in\mathbb{R}^{D\times D\times D}$ acts as a filter, selecting and emphasizing the most informative terms.

Definition 3: A mask operation $\mathbf{M} \in \mathbb{R}^{D\times D}$ is applied to randomly sample elements in $\mathbf{H}_ {i}$, and the resulting masked tensor is processed using a learnable weight tensor $\mathbf{W} \in \mathbb{R}^{D\times D\times D}$ as follows: $\tilde{\mathbf{h}}_ i$ = ($\mathbf{M} \odot\mathbf{H}_ {i}): \mathbf{W}$, where $\odot$ represents element-wise multiplication and $:$ denotes double contraction of tensors. The resulting feature $\tilde{\mathbf{h}}_ {i} \in\mathbb{R}^{1\times D}$ enriches the representation of $\mathbf{h}_ {i} \in \mathbb{R}^{D}$.

Corollary 2 (Representation Power): The full feature map $\mathbf{H}_i$ contains $D^2$ terms for a $D$-dimensional input feature vector $\mathbf{h}_i$. After masking, the effective representation space reduces by the sparsity of $\mathbf{M}$. The learnable filter $\mathbf{W}$ further narrows this down to the most critical terms.

Dear Reviewer KT6u,

Again, thanks for your constructive comments. We would like to follow up on our rebuttal to ensure that all concerns have been adequately addressed. If there are any further questions or points that need discussion, we will be happy to address them. Your feedback is invaluable in helping us improve our work, and we eagerly await your response.

Thank you very much for your time and consideration.

Best regards,

The Authors

Dear Reviewer KT6u:

As the author-reviewer discussion period will end soon, we would appreciate it if you could review our responses at your earliest convenience. If there are any further questions or comments, we will do our best to address them before the discussion period ends.

Thank you very much for your time and efforts. Looking forward to your response!

Sincerely,

The Authors

Dear Reviewer KT6u,

Again, thanks for your constructive comments, which are very much helpful for improving our paper. If there are any further questions or points that need discussion, we will be happy to address them. Your feedback is invaluable in helping us improve our work, and we eagerly await your response.

Moreover, we have thoroughly proofread our paper, corrected typos and grammar mistakes, and re-organized the contents to improve the clarity of the paper. We believe the presentation has been substantially improved (see revisions marked in red color). Please refer to the updated .pdf file.

Thank you very much for your time and consideration.

Best regards,

The Authors

I appreciate the authors' response. However, I still lean to rejecting this paper for the following reasons:

• As Reviewer rgEv pointed out, the comparison baselines are still outdated. Most of the methods being compared are from 2021-2023, which raises doubts about whether this truly demonstrates the superiority of the proposed method.

• The credibility of the comparisons is questionable. The newly added GeoFNO (2023) crashes in most experiments, while Transolver (2024) performs significantly worse than MGN. I'm uncertain about the value of these comparisons. It either suggests that these baseline methods are not meaningful for comparison, or the authors may have implemented the baselines incorrectly.

• While we should be cautious in our assessment, I still find that this paper lacks novelty in its contribution.

Thank you for your feedback. Your time and effort placed on reviewing our paper are greatly appreciated!

Comment: Regarding the alleged "outdated" baselines (2021-2023).

Reply: We appreciate but respectfully disagree with the comment that the baselines used in our study are "outdated". In fact, we carefully selected a range of baselines (eight in total) that are widely recognized in the field (e.g., MGN [1] (ICLR, 2021), GAT [2] (arXiv, 2017), GATv2 [3] (arXiv, 2022), FNO [4] (ICLR, 2021)) and some latest representative methods (e.g., MP-PDE [5] (ICLR, 2022), FFNO [6] (ICLR, 2023), Geo-FNO [7] (NIPS, 2024), Transolver [8] (ICML, 2024)) for comparison. Many of these have been recently published in prestigious venues and actively used as representative baseline models in recent research. Hence, we believe that the baseline comparisons are up to date in the field of spatiotemporal dynamics prediction.

Meanwhile, we have provided extensive experiments (see Tables 1-5, Appendix Tables S13-S17 in the revised paper) comparing our method with all baselines to demonstrate our model's effectiveness. If the reviewer believes that there are more recent or more appropriate baselines we should include, we sincerely welcome your suggestions and would be more than happy to include them in final version of the paper. Thanks!

Comment: Regarding the credibility of the comparisons and implementation concerns.

Reply: In fact, the underperformance of GeoFNO [7] (NIPS, 2024) and Transolver [8] (ICML, 2024) on limited training datasets likely stem from their mapping techniques between the irregular mesh and the regular grid, which rely on diverse training data. This is a quite common issue when these types of models are trained in small data regimes like the case we are specifically considering in this paper.

Moreover, we would like to draw your attention that all the baselines were implemented with the utmost care, following publicly available codebases and the instructions provided in the corresponding papers. We ensure you that these implementation are correct, where the baseline comparison codes will be released along with our codes after the peer review stage.

Our experiments (see Table 1 in the revised paper) aim to evaluate all the methods across diverse settings. The results highlight the robustness of our model in scenarios where other baselines may struggle, particularly when the training datasets are limited. More detailed information of these models have been provided in Appendix Subsection D.2, Appendix Tables S4-S12 in the revised paper.

Hope this clarifies your concern.

Comment: Regarding the novelty of our contribution.

Reply: In fact, the novelty of our proposed CeGNN model includes the introduction of two key modules in the network architecture, namely,

• (1) the two-level cell-embedded message passing mechanism, which better captures the spatial dependency;
• (2) the unique feature-enhanced (FE) block, which further learns the distinguishable features.

With extensive experiments on various spatiotemporal systems, CeGNN achieves superior performance compared with other baseline models across all benchmark datasets, as shown in Table 1 in the revised paper (Page 9). This has been thoroughly discussed and clearly demonstrated from the theoretical perspective (see Subsubsection 3.1.2, Appendix Subsection C.1, C.2, D.3 in the revised paper) and evidence of the empirical results (see Tables 1, 3 and 4, Appendix Tables S13-S17 in the revised paper). Hope this clarifies your concern.

We also recognize that novelty can sometimes be nuanced, sometimes depending on reader's personal judgement. However, our proposed network architeture, consisting of the cell-embedded message passing mechanism and the FE block, is unique and new, with superior performance clearly contributing to its novelty.

Concluding Remark: We appreciate the reviewer’s additional comments and critical evaluation. We sincerely hope to have your re-evaluation of our paper in light of our clarifications and contributions. Your possible consideration of updating the score is highly appreciated!

We look forward to your feedback!

References:

[1] Pfaff et al. Learning mesh-based simulation with graph networks. In ICLR, 2021.

[2] Veličković et al. Graph attention networks. arXiv 2017.

[3] Brody S, Alon U, Yahav E. How attentive are graph attention networks? arXiv 2022.

[4] Li et al. Fourier neural operator for parametric partial differential equations. In ICLR, 2021.

[5] Brandstetter et al. Message passing neural PDE solvers. In ICLR, 2022.

[6] Tran et al. Factorized fourier neural operators. In ICLR, 2023.

[7] Li et al. Geometry-informed neural operator for large-scale 3d pdes. In NIPS, 2024.

[8] Wu et al. Transolver: A fast transformer solver for pdes on general geometries. In ICML, 2024.

Thanks for your constructive comments!

Weakness 1a: The concern on novelty of cell-attribution.

Reply: Thanks for your comment. In fact, the novelty of our proposed CeGNN model includes the introduction of two key modules in the network architecture, namely, (1) the two-level cell-embedded message passing mechanism and (2) the unique feature-enhanced (FE) block. Given the cell-embedded mechanism that better captures the spatial dependency and the FE block that further learns the distinguishable features, CeGNN achieves superior performance compared with other baseline models across all benchmark datasets, as shown in Table 1 in the revised paper (Page 9).

However, following your suggestion, we have added many contents in the revised paper (see Subsection 3.1 on Pages 4-6 and Appendix Tables S16, S17 on Page 23).

• A primary theoretical deduction is added in Subsection 3.1.
• Two experimental results are added in Appendix Tables S16 and S17.
• More discussion have also been added in Appendix Subsubsection D.3.5.

Generally, the traditional message passing (MP) mechanism can be regarded as a refinement on a discrete space, analogous to an interpolation operation, which implies that edges are essentially interpolated from nodes. A MP mechanism introducing the cell further enhances the refinement of the discrete space (e.g., secondary refinement), thereby reducing the magnitude of discretization errors spatially, paving the way for its application in complex graph structures. Please see below Definition 1 on cell in graph and Corollary 1 on expressive power of cell-embedded MP. Therefore, we proposed a new two-level cell-embedded mechanism to process the message on graphs, which forms the key novelty of the proposed model.

Definition 1: (Cell in Graph) Let $G = (V, E)$ be a graph, where $V$ is the set of nodes $\mathbf{v}$ and $E \subseteq V \times V$ is the set of edges. A cell in $G$ is a subset of nodes $C \subseteq V$, such that the nodes in $C$ form a complete subgraph (clique) or satisfy predefined structural relationships. In particular, a $k$-cell $C_k$ in a graph $G$ contains $k+1$ nodes, where $\forall i, j \in C_k,$, $(\mathbf{v}_i, \mathbf{v}_j) \in E$, representing various structures, such as node ($k=0$), edge ($k=1$), triangle ($k=2$), tetrahedron ($k=3$), and so on.

Corollary 1: (Expressive Power) Given a graph $G$ including many $k$-cell ($k=0,1,2,\dots$), there exists a cell-based scheme that is more expressive than Weisfeiler-Lehman (WL) tests in distinguishing non-isomorphic graphs (see the proof in Appendix Subsection C.2 of the revised paper (Pages 16-17)).

Weakness 1b: Differences compared with exisiting topological graph learning methods.

Reply: Based on the references mentioned by you, we focus on discussing CCNNs [1] and MPSNs [3], since the work in [2] shares similarity with CCNNs [1]. A summary of these models is shown in the following Table D. We have added this new table to our revised paper, referring to Appendix Table S16 in the revised paper (Page 22).

Table D: Summary of CeGNN, MGN, and the works in [1] and [3].

Firstly, MGN achieves message passing through a two-level structure ($edge \rightarrow node$). On this basis, CCNNs in [1] leverages combinatorial complexes to achieve message passing through a three-level structure ($cell \rightarrow edge \rightarrow node$). Meanwhile, MPSNs in [3] performs message passing through a two-level structure ($complexes \rightarrow simplex$) on various simplicial complexes (SCs), which primarily include one simplex (node) and four types of complexes with varying adjacent simplices to enhance feature distinguishability and thus improve classification performance.

In contrast, our CeGNN model sets apart from these works and adopts a novel two-level structure ($[cell, edge] \rightarrow node$) from the spatial perspective, which is more suitable to learn the underlying spatiotemporal dynamics. Since there are only node labels for the supervised learning, the three-level message passing mechanism in [1] failed to yield performance lift which might be due to the introduction of redundant path of information passing. However, our two-level message passing sequence reduces the high coupling degree in [1] and avoids the limitation for additional predefined special complexes (e.g., some simplicial complexes) in [3]. A comparison result between two-level and three-level mechanisms is shown in the following Table E. The results in Table E demonstrate that the three-level mechanism underperforms compared with our proposed two-level mechanism in all cases. We have added the new results to our revised paper, referring to Appendix Table S17 in the revised paper (Page 23). Moreover, as shown in Table 3 in the revised paper (Page 9), we have provided an ablation study on all benchmarks to assess the contributions of the FE and CellMPNN blocks in CeGNN.

Table E: Comparison result between two-level and three-level mechanisms

References:

[1] Hajij, et al. Topological deep learning: Going beyond graph data. arXiv:2206.00606, 2022.

[2] Hajij, et al. Cell complex neural networks. arXiv:2010.00743, 2020.

[3] Bodnar, et al. Weisfeiler and lehman go topological: Message passing simplicial networks, In ICML, 2021.

We appreciate your constructive comments and suggestions. Please let us know if you have any other questions. We look forward to your feedback!

Dear Reviewer 1r7X,

Again, thanks for your constructive comments. We would like to follow up on our rebuttal to ensure that all concerns have been adequately addressed. If there are any further questions or points that need discussion, we will be happy to address them. Your feedback is invaluable in helping us improve our work, and we eagerly await your response.

Thank you very much for your time and consideration.

Best regards,

The Authors

Dear Reviewer 1r7X:

As the author-reviewer discussion period will end soon, we would appreciate it if you could review our responses at your earliest convenience. If there are any further questions or comments, we will do our best to address them before the discussion period ends.

Thank you very much for your time and efforts. Looking forward to your response!

Sincerely,

The Authors

Dear Reviewer 1r7X,

Again, thanks for your constructive comments, which are very much helpful for improving our paper. If there are any further questions or points that need discussion, we will be happy to address them. Your feedback is invaluable in helping us improve our work, and we eagerly await your response.

Moreover, we have thoroughly proofread our paper, corrected typos and grammar mistakes, and re-organized the contents to improve the clarity of the paper. We believe the presentation has been substantially improved (see revisions marked in red color). Please refer to the updated .pdf file.

Thank you very much for your time and consideration.

Best regards,

The Authors

Dear Reviewer 1r7X,

As the rebuttal period has undergone over two weeks, your silence made us anxious. We would like to follow up on our rebuttal to ensure that all your concerns have been adequately addressed. If there are any further questions or points that need discussion, we will be happy to address them. We eagerly await your response.

Thank you very much for your time and consideration.

Best regards,

The Authors

Thanks for your positive feedback and constructive comments.

Weakness: More theoretical explanation of FE block.

Reply: Great suggestion! Following your suggestion, we have added detailed explanation in the revised paper (see Section 3.1.1 on Pages 4-5 and Appendix Section C.3 on Page 17).

• We updated the FE block in Figure 2 to further explain its process and redefined Algorithm 1 (see Appendix).
• A primary theoretical deduction was updated in Subsubsection 3.1.1.
• More theoretical knowledge have also been added in Appendix Section C.3.

The FE block is inspired by interaction models in physics and mathematically hypothesized to capture nonlinear dependencies, enhancing the model's representation power.

Hypothesis. The FE block's hypothesis could be framed as: By capturing second-order interactions between latent features and applying selective filtering, the model can better represent complex structures or relationships in the data.

Physical Analogy. In physics, the outer product and second-order terms are often used to model interactions, such as stress tensors in mechanics or pairwise correlations in quantum mechanics. Here, the module could draw an analogy to systems where interactions between individual components (features) are crucial to the overall behavior.

Process Overview. In detail, it regards the node latent feature $\overline{\mathbf{h}}_ i \in\mathbb{R}^{D\times 1}$ as basis and builds a higher-order tensor feature $\mathbf{H}_ {i} \in\mathbb{R}^{D\times D}$ via an outer-product operation, e.g., $\overline{\mathbf{h}}_ {i}\otimes\overline{\mathbf{h}}_ {i}$. This process creates abundant second-order nonlinear terms to enrich the feature map. We then use a mask operation with $\mathbf{M}\in\mathbb{R}^{D\times D}$ to randomly sample these terms, filtering the appropriate information by a learnable weight tensor $\mathbf{W}\in\mathbb{R}^{D\times D\times D}$ to enhance the model's representation capacity.

Outer Product as Basis Expansion. The outer product operation $\otimes$ on the reshaped feature map $\overline{\mathbf{h}}_ {i}\in\mathbb{R}^{D\times 1}$ expands the original latent feature space into a higher-order tensor space. This expansion introduces second-order terms (e.g., $\alpha\beta$ for $\alpha, \beta \in \overline{\mathbf{h}}_{i}$), which can capture interactions between individual components of the original feature $\mathbf{h} _{i} \in\mathbb{R}^{D}$. Mathematically, the second-order tensor reads $\overline{\mathbf{h}} _{i} \otimes\overline{\mathbf{h}} _{i}$.This operation creates a richer feature map with cross-term interactions that may not be explicitly encoded in the original latent space.

Lemma 1 (Nonlinear Representation): The second-order terms $\alpha\beta$ can model nonlinear dependencies between features. This is particularly useful for capturing complex interactions that linear transformations (e.g., simple dot products) might overlook.

Definition 1: The FE block expands the latent feature $\overline{\mathbf{h}}_ i \in\mathbb{R}^{D\times 1}$ of node $i$ into a higher-order tensor space using an outer product: $\mathbf{H}_ {i}=\overline{\mathbf{h}}_ {i} \otimes\overline{\mathbf{h}}_ {i}$, where $\mathbf{H}_ {i}\in\mathbb{R}^{D\times D}$ is a higher-order feature map.

Regularization via Masking. Masking introduces sparsity in $\mathbf{H}_ {i}$, reducing overfitting. If $\mathbf{M}_ {jk}$ is selected, $\mathbf{M}_ {jk} =1$. Otherwise, $\mathbf{M}_ {jk} =0$. Here $j$ and $k$ index the $M \in\mathbb{R}^{D\times D}$ components.

Theoretically, this operation serves two purposes: (1) Reducing computational complexity by randomly sampling terms from the higher-order feature space. (2) Regularizing the model by introducing sparsity, which can prevent overfitting in high-dimensional spaces.

Learnable Filtering. The learnable weight tensor $\mathbf{W}\in\mathbb{R}^{D\times D\times D}$ acts as a filter, selecting and emphasizing the most informative terms.

Definition 2: A mask operation $\mathbf{M} \in \mathbb{R}^{D\times D}$ is applied to randomly sample elements in $\mathbf{H}_ {i}$, and the resulting masked tensor is processed using a learnable weight tensor $\mathbf{W} \in \mathbb{R}^{D\times D\times D}$ as follows: $\tilde{\mathbf{h}}_ i$ = ($\mathbf{M} \odot\mathbf{H}_ {i}): \mathbf{W}$, where $\odot$ represents element-wise multiplication and $:$ denotes double contraction of tensors. The resulting feature $\tilde{\mathbf{h}}_ {i} \in\mathbb{R}^{1\times D}$ enriches the representation of $\mathbf{h}_ {i} \in \mathbb{R}^{D}$.

Corollary 1 (Representation Power): The full feature map $\mathbf{H}_i$ contains $D^2$ terms for a $D$-dimensional input feature vector $\mathbf{h}_i$. After masking, the effective representation space reduces by the sparsity of $\mathbf{M}$. The learnable filter $\mathbf{W}$ further narrows this down to the most critical terms.

Thank you very much!

Thanks for your constructive comments!

Weakness 1: Add more recent works as baselines.

Reply: Following your suggestion, we have cited several important studies in recent years [1], especially those published after 2022 (see Section 2 on Page 4) and added detailed experimental results of the baseline comparison (see Subsection 4.3 on Page 7, Table 1 on Page 9, Appendix Subsection D.2 on Page 19, Appendix Tables S10-S12 on Pages 20-21) in the revised paper.

Concretely, we have considered some new baselines (e.g., FFNO, Geo-FNO, Transolver) to compare with our CeGNN model on our all cases. A summary of the results of these models is shown in the following Table A. In particular, we found that the Geo-FNO [2] and Transolver [4] perform poorly on the limited training dataset, which is significantly smaller, by at least an order of magnitude less, than the hundreds or thousands of datasets typically used for such models. We have added the new results to our revised paper, referring to Table 1.

Table A: RMSE metrics for CeGNN and other models.

References:

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

[2] Li, et al. A. Fourier neural operator with learned deformations for pdes on general geometries. arXiv:2207.05209, 2022.

[3] Tran, et al. Factorized fourier neural operators. In ICLR, 2023.

[4] Wu, et al. Transolver: A fast transformer solver for pdes on general geometries. In ICML, 2024 (Spotlight).

Weakness 2a: Is CeGNN an extension of MP-PDE? Whether the cell features enhance the performance of MP-PDE or not.

Reply: Thanks for your comment. In fact, our method is not a straightforward extension of MP-PDE.

Briefly, the key innovation of MGN-based method, MP-PDE, lies in replacing the MLP in MGN's decoder with a 1D-CNN block to aid autoregressive temporal marching. Furthermore, the training strategy, aka, the pushforward trick and temporal bundling trick, also leads to its performance improvement.

However, the innovation of our CeGNN model is the two-level cell-embedded message passing mechanism and the unique feature-enhanced (FE) block. Moreover, all our experiments are trained with an end-to-end one-step training strategy, rather than the pushforward trick and temporal bundling trick used in MP-PDE.

Given the cell feature that better captures the spatial dependency and the FE block that further learns the distinguishable feature, CeGNN achieves superior performance compared with other baselines, significantly reducing the prediction errors on all benchmarks, as shown in Table 1 in the revised paper and Table A shown above.

Weakness 2b: How much improvement could be observed if this feature was integrated into MP-PDE?

Reply: Following your suggestion, we have conducted experiments incorporating the cell feature in MGN and MP-PDE models and added detailed results in the revised paper (see Subsection 4.3 on Page 7, Appendix Subsubsection D.3.3 on Pages 21-22, Appendix Table S14 on Page 22).

Specifically, we consider MP-PDE integrating the cell feature as our baseline, listed as "MP-PDE + Cell" in the following Table B. Meanwhile, the results of MGN integrating the cell feature have been also listed as "MGN + Cell". The results in Table B show that the cell feature improves the performance of MP-PDE, but the "MP-PDE + Cell" still underperforms our model. More importantly, the results further explain why we integrate the cell feature into MGN rather than MP-PDE. We have added the new results to our revised paper (referring to Appendix Table S14).

Table B: RMSE metrics for CeGNN, MGN, MGN+Cell, MP-PDE, and MP-PDE + Cell.

Weakness 3: Dataset source; training dataset size; the experiments on the self-generated training data are not convincing.

Reply: Thank you for your comments.

The source of training data.

In fact, all the data used in this paper are from the publicly available datasets. Here, the datasets of four synthetic PDE systems are from the work in [5], while the real-world Black Sea dataset is taken from [6].

References:

[5] Rao, et al. Encoding physics to learn reaction–diffusion processes. Nature Machine Intelligence, 2023, 5(7): 765-779.

[6] Lima, et al. Black Sea Physical Reanalysis (CMEMS BS-Currents). Copernicus Monitoring Environment Marine Service (CMEMS), 2020.

The training dataset size.

In fact, we have provided the trajectory number for training, validation, and testing, described by the form like (a/b/c) in Appendix Table S2 (please see Page 18 in the paper). For example, the BS training/validation/testing dataset size is (20/2/2).

Experiments on the self-generated training data are not convincing.

In fact, we have experimented not only on the synthetic PDE systems, but also the real-world data, which are from the publically available datasets used in previously published papers. All the results in Table 1 in the revised paper (Page 9) have shown that CeGNN achieves superior performance compared with other baseline models. Hope this clarifies your concern.

Weakness 4: Long-term simulations with one-step prediction.

Reply: Thanks for your comment. All experiment tasks are tested on long-term prediction rather than one-step prediction. Once the model is trained (based on one-step prediction), we employ multi-step rollout for long-term prediction. In fact, we have described it in Subsection 4.2 (see Lines 339-342 on Page 7), shown as follows.

• "... we mainly focus on predicting much longer time steps with lower error and attempt to achieve better generalization ability ... utilize the one-step training strategy ..."

More rigorously, our research task is to predict all the next $n$ steps given a random initial condition $\mathbf{u}_ {0}$ taking the autoregressive form $\mathbf{u}_ {0} \rightarrow \mathcal{F}(\mathbf{u}_ {0}) \rightarrow \hat{\mathbf{u}}_ {1} \rightarrow \mathcal{F}(\hat{\mathbf{u}}_ {1}) \rightarrow \hat{\mathbf{u}}_ {2} \rightarrow \cdots \rightarrow \hat{\mathbf{u}}_ {n}$, where the function $\mathcal{F}$ is an unknown function learned by our model, and we set the $n$ steps as hundreds or thousands of steps in the synthetic datasets (e.g., 1000 steps in 2D Burgers, 3000 steps in 2D FN, 3000 steps in 2D GS, 3000 steps in 3D GS), and 20 steps in the real-world dataset.

Weakness 5: The significance of cell feature; Can its significance be verified or be measured?

Reply: Good question! In fact, the CeGNN's innovation includes the two-level cell-embedded message passing mechanism and the unique feature-enhanced (FE) block. Given the cell-embedded mechanism that better captures the spatial dependency and the FE block that further learns the distinguishable features, CeGNN achieves superior performance compared with other baseline models across all benchmark datasets, as shown in Table 1 in the revised paper (Page 9).

Following your suggestion, we added explanation in the revised paper (Subsection 3.1 on Pages 4-6 and Appendix Section C on Pages 16-17).

• A primary theoretical deduction in Subsection 3.1.
• More theoretical preliminaries in Appendix Section C.

Generally, the traditional message passing (MP) mechanism can be regarded as a refinement on a discrete space, analogous to an interpolation operation, which implies that edges are essentially interpolated from nodes. A MP mechanism introducing the cell further enhances the refinement of the discrete space (secondary refinement), thereby reducing the magnitude of discretization errors spatially, paving the way for its application in complex graph structures.

Definition 1 (Cell in Graph): Let $G=(V, E)$ be a graph, where $V$ is the set of nodes $\mathbf{v}$ and $E\subseteq V\times V$ is the set of edges. A cell in $G$ is a subset of nodes $C \subseteq V$, such that the nodes in $C$ form a complete subgraph (clique) or satisfy predefined structural relationships. In particular, a $k$-cell $C_k$ in a graph $G$ contains $k+1$ nodes, where $\forall i, j \in C_k,$, $(\mathbf{v}_i, \mathbf{v}_j) \in E$, representing various structures, such as node ($k=0$), edge ($k=1$), triangle ($k=2$), tetrahedron ($k=3$), and so on.

Corollary 1: (Expressive Power) Given a graph $G$ including many $k$-cell ($k=0,1,2,\dots$), there exists a cell-based scheme that is more expressive than Weisfeiler-Lehman (WL) tests in distinguishing non-isomorphic graphs (see the proof in Appendix Subsection C.2 of the revised paper (Pages 16-17)).

Hence, we proposed a new two-level cell-embedded mechanism to process the message on graphs.

Weakness 6: The benefit of adding position awareness; the performance comparison of replacing the distance to the cell center with the distance to the nearest PDE boundary for each point.

Reply: Thanks for your comment. Following your suggestion, we tested other two variant models, namely, "CeGNN w/o Cell Pos." and "CeGNN w Cell (Pos. to B)", with the results reported in the following Table C. Here, "CeGNN w/o Cell Pos." represents the cell initial feature without the position awareness, and "CeGNN w Cell (Pos. to B)" replaces the distance to the cell center with the distance to the nearest PDE boundary. The results in Table C show that the performance improvement of CeGNN is not only due to adding position awareness into the cell initial feature, but also the secondary refinement of the discrete space. We have added the new results to our revised paper (referring to Appendix Table S15 on Page 22 in the revised paper).

Table C: RMSE metrics for CeGNN and other two cases.

Overall, we appreciate your constructive comments and suggestions. Please let us know if you have any other questions. We look forward to your feedback!

Most of my concern is addressed. But I observe that the time step setting is different between different datasets. Why they have such a large gap. Does it can be attributed to the different time interval setting in different setting? For example, 1000 steps in your 2D Burgers setting takes only a few seconds.

Dear Reviewer rgEv:

As the author-reviewer discussion period will end soon, we would appreciate it if you could review our responses at your earliest convenience. If there are any further questions or comments, we will do our best to address them before the discussion period ends.

Thank you very much for your time and efforts. Looking forward to your response!

Sincerely,

The Authors

This is an excellent paper, particularly after addressing all the reviewers' comments. It is well-written, well-organized, and reflects a high level of effort and attention to detail. I firmly believe the authors deserve constructive communication or feedback to acknowledge their dedication.