Discovering Message Passing Hierarchies for Mesh-Based Physics Simulation

Q3. About the AMP layers.

We use AMP over traditional attention mechanisms for its better performance in dynamic systems. First, we have compared the results of DHMP with HCMT, which employs an attention-based GNN architecture. As shown in Table 2, DHMP outperforms HCMT in all cases.

Second, we indeed considered using scaled-dot product attention instead of the AMP module by replacing the original edge update function with a cross-attention that integrates both edge features and neighboring node features. Specifically, the edge feature is updated as $\mathbf{\hat{e}}\_{aggr} = \operatorname{CrossAttention}(Q = \mathbf{v}\_i, \, K = \{\hat{\mathbf{e}}\_{ij}\}, ,V = \{\hat{\mathbf{e}}\_{ij}\})$, where the edge feature is refined by attending to node features. The cross-attention scores are then reused for inter-level information propagation. The results are as follows:

The results demonstrate that DHMP with AMP outperforms the cross-attention variant in overall accuracy. This is likely because the edge features $\hat{\mathbf{e}}\_{ij}$ inherently incorporate information from the associated nodes $\mathbf{v}\_i$, making it unnecessary to explicitly include $\mathbf{v}\_i$ in the computation, which would introduce redundancy.

Q4. Increased computational complexity by using AMP instead of GraphNetBlock.

The increased computational complexity when using AMP instead of a regular GraphNetBlock comes primarily from the addition operation of importance weights and the softmax normalization.

For both GraphNetBlock and AMP, given the $\mathbf{e}\_{ij}$ with constant dimension $c$ and $\hat{\mathbf{e}}\_{ij}$ and $\mathbf{v}\_{i}$ with dimension $D$, the edge update and node update functions have computational complexities of $O(|E| \cdot (cD + D^2))$ and $O(|E| \cdot D + |V| \cdot D^2)$. AMP computes the importance weight by $\phi^w$, which incurs an additional complexity of $O(|E| \cdot (cD + D^2))$. Furthermore, the softmax operation results in an additional complexity of $O(|E| \cdot D)$.

Although the computation of the importance weights doubles the computational cost per edge compared to the regular GraphNetBlock, the overall complexity still remains $O(|E|D + |V|D^2)$.

In the following table, we replace the GraphNetBlock in MGN with AMP to illustrate the increased computational cost (data: CylinderFlow).

Q5&6. Time & memory cost of the construction of K-hop edges. Appendix D2: The training/inference time and vRAM usage for the different K-values.

In the revised Table 9 in Appendix D.2, we present comparisons of training/inference time and vRAM usage for edge construction of different $K$-values on the FlagSimple dataset. As $K$ increases, both the training and inference times gradually rise due to the enhanced edge construction, which involves more computational overheads.

Changes and additional comments.

Thank you for your careful checking and constructive feedback. We have revised our paper accordingly in L108, L116, L152, Table 2, L761, and L807. Additionally, we have included reference [4] in the related work and presented more details of the model architectures in Appendix C.

We appreciate the reviewer’s detailed comments and hope that the following responses can adequately address the reviewer’s concerns. We are more than willing to address any further questions you may have.

W1. Experiment details and performance discrepancy.

Please refer to our General Response.

W2. Missing comparison with [2].

We thank the reviewer for bringing this paper to our attention. First, we would like to kinldly remind the reviewer that, according to the ICLR reviewer guidelines: "authors are not expected to compare with work published within four months of the paper deadline". Since [2] was uploaded to arXiv just two weeks before the deadline and has not yet been officially published, it should be considered concurrent with our work.

Nevertheless, we recognize the value of such a comparison and outline the key differences from [2] below. We have also included the discussion in Appendix J.

• First, DHMP assigns different importance weights to the updated edge features, highlighting the anisotropic influence of these edges, while Garnier et al. primarily focus on attention-based interactions between updated node features.
• Second, DHMP constructs dynamic hierarchies based on the differentiable node selection achieved through AMP's predictions, while Garnier et al. introduce additional self-attention blocks to retain the top $k$ nodes.

W3. l202-3-4: I think that statement is wrong. GAT networks for example do use edge features during the attention computation. Similarly, even with absolute node positions in the node features, why would an attention-based method not be able to compute the differences between those coordinates?

Sorry for the confusion. We have revised the manuscript to clarify this point.

Some attention-based GNNs, such as HCMT (Han et al.), do incorporate edge features when computing attention weights. However, in these methods, the resulting weights are typically applied to aggregate node features. In contrast, our proposed AMP directly applies the predicted importance weights to the edge features themselves and, more importantly, further exploits these weights in inter-level feature propagation.

Furthermore, while attention-based methods for physical simulations (e.g., HCMT and GMR-Transformer-GMUS) rely on transformer architectures with quadratic time complexity, AMP maintains a computational complexity comparable to that of a standard GraphNetBlock. This makes AMP more computationally efficient and effective for physical simulations.

Q1. l69-78: Can a parallel be made with [3] where each hop gets processed by a different MLP?

Yes, indeed! We have included this result in Appendix D.5 in the revised paper. We separate the $\mathcal{E}_{l}$ and $K$-hop enhanced edges into distinct variables to align with the approach in [3]. We replace the AMP block in DHMP with the KP-AMP block. We conduct experiments on CylinderFlow and present the results below. We can see that incorporating distinct MLPs for each hop helps to capture varying structural information at different scales. This could be a promising direction for further improvement upon DHMP.

Q2. l74/75: Where is this computational efficiency demonstrated?

Table 13 in Appendix G presents the training cost and inference time for baseline models and DHMP. The results indicate that HCMT (Yu et al.), which incorporates an attention mechanism in its message passing, requires significantly more training hours to reach convergence and exhibits a substantially higher inference time, demonstrating its greater computational overhead. Specifically, HCMT requires 1.5x to 3.9x training cost and 1.8x to 3.7× inference time compared to DHMP.

Thank you for your valuable comments, and we hope that the added results and the clarification below can address your concerns. For any further questions, please don’t hesitate to contact us.

W1. Novelty of multi-scale graph neural networks with differentiable node selection.

Thank you for bringing our attention to this related work. We have included proper citations and detailed comparisons in the revised version (Sec 5 and Appendix J). Below, we outline the key differences between MGVAE and our work.

• Usage of Gumbel-Softmax operation. The Gumbel-Softmax operation is used differently in MGVAE, where it partitions the graph into discrete clusters with a specified $K$ value at each resolution level. This approach may not be suitable for large graphs, as choosing appropriate and high $K$ values can be challenging. In contrast, DHMP uses Gumbel-Softmax for differentiable node selection without requiring hyperparameter selection.
• Edge construction at coarse levels. While MGVAE constructs edges at coarser levels by aggregating adjacency information within clusters, DHMP constructs edges using a $K$-enhanced edge set that captures long-range dependencies.
• Intra/Inter-level feature propagation. DHMP employs anisotropic message passing to aggregate neighboring features in a directionally non-uniform manner, reusing importance weights to enable efficient inter-level information propagation. However, MGVAE primarily relies on GraphNetBlocks for message passing, which builds a bottom-up aggregation and top-down decoding strategy, pooling features at each level and projecting them across the hierarchy in a variational manner.

Below, we highlight the key differences between the two methods:

W2. Complexity analysis.

The asymptotic complexity of our approach is $O(|E|d + |V|d^2)$, which is consistent with standard GNNs. For baselines, the message passing GraphNetBlock with MLP implementation introduces a time complexity $O(|E|d + |V|d^2)$, where $|E|$ is the number of edges and $|V|$ is the number of nodes, $d$ is the feature dimension. Methods like MGN process message passing on high-resolution meshes, resulting in higher computational overhead. In contrast, hierarchical methods like BSMS-GNN reduce computation since coarser graph levels have fewer edges and nodes. In our proposed DHMP, the additional operations primarily come from the dynamic graph construction and anisotropic message passing, with detailed complexity as follows.

• Dynamic Graph Construction: (1) Probability-based node selection involves computing probabilities for all nodes, resulting in a complexity of $O(|V|)$. (2) After node selection, edges are pooled to form a coarsened graph. This operation involves evaluating and updating edges for the selected nodes, with a complexity of $O(|E|)$.
• Anisotropic Message Passing: AMP introduces an overhead for predicting directional weights on edges. We compute the directional weight and perform weighted aggregation, maintaining a complexity of $O(|E|d + |V|d^2)$.

Overall, the asymptotic complexity of DHMP is $O(|E|d + |V|d^2)$, which is consistent with standard GNNs. The computation of the importance weights doubles the computational cost per edge compared to the regular GraphNetBlock. Therefore, the computational overhead is acceptable.

W3. Stability of Differentiable Node Selection (DiffSELECT) and sensitivity to temperature hyper-parameter \tau.

(1) The influence of Gumbel-softmax sampling on performance stability.

We showcase the stability of our model in three aspects:

• Test stability: As shown in Table 11 in Appendix E, DHMP provides stable results under independent tests with Gumbel-Softmax sampling.
• Training stability: In Table 12 in the revised Appendix F, we present a full comparison between our model and the baseline models, calculated across three random seeds. The shown standard deviations highlight the advantage of DHMP in maintaining stable training.
• Robustness to out-of-distribution systems: To further evaluate stability in highly dynamic systems, we have evaluated DHMP on datasets with significantly higher velocity norms and variances (Table 5 in Sec 4.4). As shown below, the standard deviation of prediction errors remains minimal, highlighting that once trained, DHMP consistently produces stable predictions under Gumbel-softmax sampling, even in highly dynamic scenarios.

(2) Sensitivity to the Gumbel-Softmax temperature during training.

We acknowledge the sensitivity of the Gumbel-Softmax to temperature ($\tau$). In our experimental setting, we employed temperature annealing during training, gradually decreasing the temperature from 5 to 0.1 with exponential decay on all datasets, as detailed in Appendix B (lines 817-819). This approach balances exploration in early training with progressively refined node selection afterwards, enhancing both stability and performance while reducing the need for hyperparameter tuning [1-2]. Table 12 in Appendix F shows the results of the mean and standard deviation calculated over three random training seeds under this setting.

Additionally, we have experimented with fixed temperatures ($\tau = 0.1, 0.3, 0.5$) on the CylinderFlow dataset. The results are shown below:

[1] Categorical reparameterization with gumbel-softmax.

[2] The concrete distribution: A continuous relaxation of discrete random variables.

Dear Reviewer, Do you mind letting the authors know if their rebuttal has addressed your concerns and questions? Thanks! -AC

Dear rDem,

Thank you once again for your review of our work. As the discussion period is approaching its end, we would be grateful if you could reply whether our responses have addressed your questions and concerns. Your time and effort are greatly valued.

We appreciate the reviewer’s detailed comments and hope that the following responses can adequately address the reviewer’s concerns. We are more than willing to address any further questions you may have during the discussion period.

W1 & Q1: Nodes distributed in uninformative regions; Q2: Node distribution over trajectories.

(1) "Uninformative" regions.

As a data-driven neural network, DHMP does not have access to physical priors, such as knowledge of where turbulence or laminar flow occurs. As a result, it may retain nodes in regions that could be perceived as "uninformative" from a physical standpoint.

However, it's important to note that these uninformative regions, particularly those with laminar flow, are not necessarily trivial or easy. As illustrated in the error map at the bottom of Figure 3, baseline models like MGN and BSMS-GNN show higher errors in these regions, indicating that they present challenges for these models as well. In these challenging regions (even with laminar flow), excluding interactions from these areas entirely could hinder the model’s ability to learn a more comprehensive representation of the system dynamics. These regions can still influence other parts of the system. In other words, retaining these nodes in top hierarchies is crucial for a more complete understanding of the overall dynamics and long-range node relations.

(2) Node distribution across the entire sequence.

We have visualized the evolution of the graph structure across the entire sequence in the GIF files uploaded in the revised supplementary material. The results demonstrate that nodes are typically distributed in regions with larger temporal differences, supporting our observation that the model prioritizes areas with more dramatic changes in fluid motion.

W2&Q3. About the number of hierarchies. How is the number of hierarchy selected? ... based on the problem size/difficulty?

Thank you for the insightful question! The number of hierarchy levels in DHMP is generally determined based on the problem size, roughly $O(\log |V|)$. This choice is also consistent with static hierarchy methods such as BSMS-GNN.

We acknowledge that the hierarchy parameter may impact model performance differently, so we conducted a comparison on the CylinderFlow dataset to analyze DHMP and BSMS-GNN under various hierarchy parameter settings. The results are as follows and have been included in Figure 9 in Appendix D4.

The results indicate that DHMP outperforms BSMS-GNN at all hierarchy levels. Both models benefit from increased hierarchy depth up to level 7, suggesting deeper levels enhance accuracy by capturing more complex interactions. However, a slight performance drop at level 9 hints at possible overfitting. Overall, the dynamic hierarchical models in DHMP consistently outperform their counterparts with static hierarchies in BSMS-GNN.

W3 & Q4. Discrepancy in model performance. Experiment setup and implemantation details.

Please refer to our General Response.

Thank you for your comments. We have made every effort to answer each of your questions, and hope that our response can address your concerns regarding this paper.

W1. The proposed anisotropic message passing differs only in small implementation details from graph attention.

We would like to clarify the novelty and contributions of our method, particularly addressing the role of AMP within the broader framework of dynamic hierarchy construction from the following aspects.

• Dynamic hierarchy construction. Firstly, we would like to emphasize primary contribution lies in constructing dynamic hierarchies that adapt to evolving physical dynamics. In this hierarchy-building process, AMP serves two key functions. First, it aggregates neighboring features to the central node in a directionally non-uniform manner. Second, it utilizes the importance weights predicted by AMP to enable efficient inter-level information propagation. The latter aspect has been unintentionally overlooked by the reviewer; however, our ablation study (Figure 4, lines 429–454) demonstrates its benefit on model performance.
• Differences between AMP and graph attention. Unlike traditional graph attention, AMP incorporates nodes' relative positions directly into edge features, capturing essential spatial relationships that drive accurate modeling of physical systems where spatial orientation greatly affects physical dynamics. Graph attention computes attention scores between nodes, while AMP directly applies the predicted importance weights to edge features. By using anisotropic interactions through edge features that embed geometric information, our approach enhances generalizability and robustness.
• Experimental results. We compared DHMP with HCMT, an attention-based baseline. As shown in Table 2 and Table 13, DHMP consistently outperforms HCMT in both performance and efficiency (especially training cost and inference time). Additionally, we explored replacing AMP with a cross-attention mechanism that integrates both edge and neighboring node features while reusing attention scores for inter-level feature propagation. This modification resulted in reduced performance, as detailed in our response to Reviewer 96Nr’s Q3.

Overall, while AMP is a key component of our dynamic hierarchy framework, it is not the sole primary contribution of our method. We hope these clarifications further highlight the effectiveness and novelty of our approach.

W2 & Q1. Differences in results.

Please refer to our General Response.

The authors' main contributions are not written separately.

Thank you for your suggestion. We have revised the introduction section and listed the contributions of our work. Please feel free to share any further suggestions for improving the paper.

Dear Reviewer, Do you mind letting the authors know if their rebuttal has addressed your concerns and questions? Thanks! -AC