Learnable Stability-Aware Unstructured Grid Coarsening Using Graph Neural Networks for Accelerated Physics Simulations

Authors: Alexander Ryabov, Sergei Shumilin, Viacheslav Naumov, Nikolay Yavich, Sayan Ranu, N M Anoop Krishnan, Evgeny Burnaev, Vladimir Vanovskiy
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
Keywords: Graph Neural Networks, differentiable solvers, numerical modelling, grid coarsening, upscaling

Abstract:
Efficient simulations of complex physical systems described by partial differential equations (PDE) require computational methods that can reduce the resource demands without sacrificing the accuracy. Traditionally, this is achieved by ``upscaling'' the simulation grids or by aggregating cells based on a priori information. Here, we introduce a novel framework based on graph neural networks (GNN) for learnable self-supervised differentiable coarsening of unstructured computational grids. We leverage graph-based representation of the physical system and offer a graph coarsening method which preserves the underlying physical properties together with the stability of the chosen numerical scheme. This is achieved by minimizing the error between the output of the simulations using coarsened and original graph. We demonstrate the approach on several example differential equations, modeling sub-surface flow and wave propagation. We demonstrate that the model exhibits ability to maintain high fidelity in simulation outputs even after 95\% reduction on the nodes, significantly reducing computational overhead. We also show that the model exhibits generalizability to unseen scenarios, thereby outperforming the baselines. Thus, the developed approach demonstrates the ability to accelerate simulation without comprising accuracy and hence has potential for accelerating physical simulations in various domains.

Decision: reject
Reviewer scores: 3, 5, 5, 5

Reviews and discussion:
<thread>
creation date:
1730360146593

thread_type:
official_review

<thread_object>
creator: reviewer_vcgh

<content>
summary:
The authors propose a novel framework that leverages graph neural networks (GNNs) to facilitate a self-supervised learnable coarsening process. This GNN-based approach creates a generalizable coarsening strategy. This adaptability allows it to function across different simulation setups without requiring grid-specific re-optimization.

strengths:
The introduction of a self-supervised GNN-based coarsening framework represents a significant advancement over traditional methods that require specific optimizations for each grid.

weaknesses:
There are missing implementation details. See questions.

questions:
1. Are alternative remeshing strategies, such as Delaunay triangulation, considered, and if so, how do their performances compare?
2. Is your proposed method adaptable to non-fixed mesh topology systems, such as Lagrangian systems like the Deforming Plate model discussed in MeshGraphNets[1]? 
3. Can you illustrate the meaning of colorbar in Figure 8 of Appendix A?
4. What is the input quantities of the GCN for each task?
5. It would be better if the authors could compare the clusters generated by the proposed method with the clusters of baselines.
6. It would strengthen the evaluation of your method to include comparisons with additional reduced-order modeling baselines, such as UPT[2] and Transolver[3].

[1] Learning Mesh-Based Simulation with Graph Networks.

[2] Universal Physics Transformers: A Framework For Efficiently Scaling Neural Operators.

[3] Transolver: A Fast Transformer Solver for PDEs on General Geometries.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730412832581

thread_type:
official_review

<thread_object>
creator: reviewer_wkfs

<content>
summary:
This paper proposes a Graph Neural Network (GNN) based approach for automated mesh coarsening for PDE solutions. The authors extend previous work on self-supervised coarsening of unstructured grids by introducing a learnable method that aims to generalize across different scenarios (eg, varying mesh resolutions and source/sink terms) without requiring new runtime optimization.

While the core idea of learnable mesh coarsening is interesting, and the authors claim to have benefits such as no need to redo optimization on the fly, there are several fundamental concerns:

1. Insufficient Performance Advantages
   - The method demonstrates worse accuracy compared to baseline in 1 out of 2 test cases; The authors should at least justify why this accuracy trade-off is acceptable in practical applications

2. Limited Generalization Capabilities
   - Experimental results show significant performance deterioration when:
     a. Applied to meshes of different resolutions
     b. Tested with additional source terms
   - This contradicts the paper's central claim of improved generalization ability

3. Limited Computational Efficiency Improvement
   - The time comparison between the proposed method and baseline approaches shows very limited advantages (eg, from 22s to 5s, both are at similar magnitudes)
   - Let along the training cost ahead of time, which makes the gain even more questionable.

Overall:
I recommend rejection of this paper. There are some additional minor suggestions on the clarity of writing, if the authors would like to improve their manuscript. But without experiments proving significant benefits, I would not consider chaning my decision.

strengths:
- An interesting trial to extend an on-the-fly optimization problem to be learned by GNN; so that it can save time in the future deployments.

weaknesses:
- See overview.

questions:
Other minor questions on clarity and writing (authors should add these details to their manuscript in addition to answering here):
- How does the GNN decides which nodes to keep in corsened mesh? Top-K?
- Once coarsened nodes are obtained, how to get the new mesh? The current method only mentions triangular mesh, what is your proposal to adapt in these cases where the mesh is 3D or no longer simplex?
- How many timesteps do you keep the gradients, 1, eg only calculating the gradients for stepping from n to n+1, or multi-step; if multi-step, how many?
- You mentioned implicit method can improve results. But where are the results to support this claim?
- The simulation is required to be differentiable. This easily achieved in explicit scheme, but challenging to general PDEs solution schemes, eg, WENO for nonlinear advection terms. What is your proposal to these extensions?
- In sec 3.3., you mentioned Sij is learnable by NN. What is the NN that outputs Sij? as Sij depends on the topology, ie the graph G, it should be output via some GNN strucutre? Otherwise it's not possible to extend to other mesh structure.
- Page 5, you mentioned stability loss (1st appearence of that term in main text) without simply introducing its meaning, which hinder the reading flow of readers.
- Can you mentioned the method to deal with cases when sensor locations do not align with the grid points? Do you use Barycentric interpolants?
- Eq. 9 does not seem to be connected with Eq. 8 as the former actually regularize the number of coarsened nodes, n, it seems to have nothing to do with eg CFL or criteria for stability. Consider improving the writing here.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730611055620

thread_type:
official_review

<thread_object>
creator: reviewer_wvwb

<content>
summary:
Aiming at accelerating the numerical PDE solving-based simulation, the authors propose a method that utilizes graph neural networks to generate a coarsened mesh such that the solver can generate solutions quickly. The loss terms are designed such that the solution is stable and accurate. The proposed framework is applied to the cases of (simplified) wave propagation and (simplified) subsurface flow.

strengths:
1. Automatic optimization of the grid, saving human efforts.
2. Designed loss term to soft-enforce stability of the simulation on the coarsened grid.
3. Results are comparable with recent methods based on global optimization.

weaknesses:
A few issues can be observed, listed below in the order of importance:

1. Whilst the solution from the coarsened grid are accurate at the node points, the super-coarsened grid inevitably lead to another source error since quite a great proportion of applications would require the field solution rather than discrete points, which would necessitate interpolations from the solution points. Such error is largely ignored when the authors are trying to coarsen the grid as much as possible. In clear words, how much error would one observe when one is going to plot out the whole simulated pressure field by interpolating from the simulation results? It is easy to imagine that the grid in Fig. 8 would give a decent interpolated field, while the coarsened grid in Fig. 9 would probably not. The authors did not consider those errors in the proposed grid coarsening algorithm, which poses a significant limitation to its application.

2. The method is applied to very simple cases. As discussed by the authors themselves, applications in more complex scenarios would be great to have in the future.

questions:
1. The largest question I have is that, since a graph neural network is already employed, why not further the effort and directly generate the solution with the graph neural network (by increasing the current one layer network to a slightly larger one)?
2. Table 1 & 2, it is reported that the method is actually slower and more memory intensive than the baseline, while results are comparable or only slightly better???
</content>
</thread_object>

</thread>
<thread>
creation date:
1730718984989

thread_type:
official_review

<thread_object>
creator: reviewer_hygj

<content>
summary:
The work proposes a method for mesh coarsening, which may facilitate faster physical simulations. The coarsening is done using self-supervised GNN to apply to unseen settings. They also investigated the effect of stability loss, which enhances the stability of training and prediction, and a differentiable implicit solver for stable prediction. The numerical experiments suggest that the proposed method has a competitive performance with a baseline machine learning method.

strengths:
* The work focuses on the mesh coarsening problem, which is relatively new to the community.
* The explanation of the method is clear and easy to follow.

weaknesses:
* The novelty of the method is limited. They claimed the differentiable implicit method is their novelty; however, PyTorch supports differentiable linear solvers ( https://pytorch.org/docs/stable/generated/torch.linalg.solve.html ). In addition, the community knows that there are implicit GNNs (e.g., [Gu et al. NeurIPS 2020 https://arxiv.org/abs/2009.06211 ]) that solve implicit problems.
* The experimental evaluation is weak. The authors state that mesh coarsening is a recent approach, and there are a limited number of existing methods, but there are actually a lot of works investigating mesh coarsening (e.g., [Pffaf et al. ICLR 2021 https://arxiv.org/abs/2010.03409 ] and [Alet et al. ICML 2019 https://arxiv.org/abs/1904.09019 ]). The authors must mention these existing works and evaluate them in the numerical experiments to show the superiority of the proposed method.
* The effectiveness of the proposed method is limited. The experimental results suggest that the method has almost the same accuracy and longer computation time for training. If the reviewer did not miss, there is no evaluation of the prediction time, which is one of the most essential metrics for the mesh coarsening task. In addition, there should be a comparison regarding the speed-accuracy tradeoff with classical numerical analysis methods with varying resolutions because a simple mesh coarsening may work for the problems considered in the paper.

questions:
* What would be the potential pros and cons compared to the classical methods, e.g., adaptive mesh refinement? The mesh coarsening task would be interesting, but the reviewer is not sure if this direction is promising, given that we have plenty of optimized methods in the classical numerical analysis field.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732547327748

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Dear Reviewers

comment:
A very important point to note for many concerns is the different nature of [Shumilin et al.] baseline from usual baselines. In the [Shumilin et al.] is solved a very difficult problem of global optimization of mesh points, which is performed for every new grid. In the proposed by us approach the GNN outputs directly a solution of this optimization problem which is much more simple and general. That's why this baseline can be considered as the one we would like to be close with (so it's more about the best achievable metrics than about the metrics we can beat). So we consider a major strength of our approach that it gives results very close to the one by [Shumilin et al.].
</content>
</thread_object>

</thread>
<thread>
creation date:
1733167197189

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Comparison with other methods

comment:
Dear Reviewers,

next, we provide the comparison with other methods as some of you have asked.

**MeshGraphNets** [1]

We compared the *accuracy* of the solution obtained by our algorithm vs MeshGraphNets. The experiment setup is as follows: we generate a complex permeability field as the input of the diffusion initial and boundary value problem. The number of initial points is $N$ = 1000. We used an explicit numerical solver to generate data as long as MeshGraphNets is trained on data. The input solution matrix is $U \in R^{N, T}$, where $T$ is the number of time steps. We took a source code and adjusted it for our diffusion problem. As the input to MGN we used node embeddings as $[x_i, y_i, k_i, p^t_i]$  - coordinates, permeability, pressure at the node $i$ at time $t$. We added some noise for $k_i, p^t_i$. For edge embeddings, we used the technique from the original paper: for $e_{ij}$ we embed $ [ \Delta x. \Delta y, ||x_i - x_j|| ] $.  As the output of the graph neural network with 8 hidden layers (k-hop size) we predicted $\partial u/ \partial t$ and then integrated it to get $u_{t+1}$. The timestep is 5e-5. We calculated the roll-out of the trained model in an auto-regressive manner: $X^{t+1} = MeshGraphNets(X^{t})$. 
We trained our method on the same grid. We located a source point in the center (as in the original dataset) and two sink (measurement, interest) points. After we applied our method and got a coarsened grid we compared the solution of the rollout and our method's solution at the sink points. The important point to note: MGN and Our method shared the same graph for GNN. The following are the results of the comparison:

[Result of coarsening](https://drive.google.com/file/d/1LaUAYD3k9nE2fKe3raoas6swh7LS2iHJ/view?usp=sharing)

[Comparison at the sink point 1 with MGN](https://drive.google.com/file/d/1_rufwFls9H_UuTC18WaWFDDFHxkShREc/view?usp=sharing)

[Comparison at the sink point 2 with MGN](https://drive.google.com/file/d/1b2wP_3Rnx_nKClzKfHacYrUYPRBFKAyf/view?usp=sharing)

Our method is much more accurate at the points of interest. Another advantage of our method is that it is faster to make a "rollout" because we leverage a differentiable numerical solver on the already coarsened grid. In contrast, MGN uses a heavy trained GNN, which inevitably increases the time of inference. 

After that, we tested the *generalizability* of both methods. To do this, we added significant Gaussian noise to the initial permeability field. And performed inference with both trained models. Our method demonstrated to be much more robust to the added noise. The following are the results of the comparison:

[Comparison of generalizability the sink point 1 with MGN](https://drive.google.com/file/d/1IzPfNy2lW5fLNQzYtDCTvSoDYlPGxtpN/view?usp=sharing)

[Comparison of generalizability at the sink point 2 with MGN](https://drive.google.com/file/d/1dBwh9DhFMEUYrKN68Txpb0JRzeh0aguC/view?usp=sharing)



**FluxGNN** [2] 

We used the source code from the repository authors provided and after lots of attempts were unable to install the code of the original paper due to errors related to code dependencies. The code is irreproducible. However, we will mention this work in the camera-ready version.

**Adaptive mesh refinement**

We provided a thorough discussion on this topic supported by experiments in the replies to Reviewer hygj.

**Shumilin 2024** [3]

The closest work which we compared with in detail throughout the paper and rebuttals.

Also in the course of the rebuttal, we compared our work with some other methods. We already updated our paper with the Related work subsection. We will include the above comparisons in the camera-ready version of our paper.


[1] https://arxiv.org/pdf/2010.03409 (LEARNING MESH-BASED SIMULATION WITH GRAPH NETWORKS)

[2] https://openreview.net/forum?id=WajJf47TUi

[3] https://openreview.net/forum?id=kMBvZ40Iu9
</content>
</thread_object>

</thread>
<thread>
creation date:
1733167300387

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Final general reply

comment:
We would like to sincerely thank all the reviewers for their constructive feedback and detailed evaluations of our work. We appreciate the time and effort each of you dedicated to reviewing our manuscript. Below, we highlight the strengths noted by the reviewers and describe the changes made to address your concerns, as well as our plans for additional improvements in the camera-ready version.

- Automatic optimization of the grid with the help of GNN, learned in a self-supervised way, saving human efforts and representing a significant advancement over traditional methods (**wVwb, WkfS, vcgH**).
- Relatively novel mesh coarsening problem (**hygj**).
- Designed stability loss term  (**wVwb**).
- Clear and easy-to-follow explanation  (**hygj**).
- Results are comparable with recent methods based on costly global optimization (**wVwb**).

**Changes made in the revised version**

We have made several improvements to the manuscript in response to the reviewers' feedback. The changes address specific concerns raised and aim to strengthen the clarity and comprehensiveness of the paper.

- Improved formulation of the third contribution as recommended by Reviewer **hygj** (**Introduction**).
- Added "Related Work" paragraph comparing our approach with recent works such as Graph Element Networks and MeshGraphNets. This was recommended by Reviewers **hygj** and **vcgH** (**Section 2.2**).
- Improved the description of the neural network architecture used for coarsening, as recommended by Reviewer **WkfS** (**Section 3.2**).
- Explained the rationale for preserving key points as recommended by Reviewer **WkfS**. Additionally, in Appendix C, we described how to handle cases where sensor locations do not align with grid points (**Section 3.2** and new **Appendix C**).
- Introduced the meaning of stability loss upon its first mention and added a link to its full definition, as recommended by Reviewer WkfS (Section 3.2).
- Described how to adapt our framework to 3D meshes or non-simplex grids, as recommended by Reviewer **WkfS** (**Section 3.2** and new **Appendix B**).
- Clarified the multi-step optimization process and stability loss function, as recommended by Reviewer **WkfS** (**Section 3.4**).
- Added a baseline METIS and incorporated new experiments, as discussed with Reviewer **hygj** (Sections **4.1**, **4.2**, **4.3**, and new **Appendix K**).
- Included experiments for the Proxy model scenario, see discussions with Reviewer **wVwb** (**Section 4.2** and new **Appendix J**).
- Added experiments with more complex permeability scenarios (log-uniform permeability), as recommended by Reviewer **wVwb** (**Section 4.5**).
- Provided motivation for using implicit solvers, as recommended by Reviewer **WkfS** (**Appendix M**).
- Incorporated a better set of hyperparameters, as discussed in our response to Reviewer **hygj** (**Appendices I.1 and I.2**).
- Moved time and memory tests to **Appendix H** for better clarity.
- Moved the formal definition of Graph Coarsening for physics-based simulations to **Appendix A**.

**Additional experiments for the camera-ready version**

We are planning further improvements for the camera-ready version, which include:

- Adding experiments comparing our algorithm with the baseline **MeshGraphNets**.
- Mentioning other works such as Universal Physics Transformers, Transolver and FluxGNN.
- Incorporating the remaining textual improvement suggestions by Reviewer **vcgH** into the manuscript.

Once again, we sincerely thank all reviewers for their thoughtful feedback, which has significantly improved the quality of our work. We are especially grateful to **wVwb** and **WkfS** for raising their scores, and we hope all reviewers will consider **raising their scores further**, taking into account the extensive improvements made to the paper and the additional experiments included.

Your insights have been invaluable, and we look forward to further refining our work in the camera-ready version.

Best regards,

The authors
</content>
</thread_object>

</thread>


Metareview:
metareview:
The work proposes a GNN based method for coarsening the mesh used in a physical simulation, promoting stability and accuracy of the solution. This can speed up traditional numerical methods as much fewer computational nodes can be used to a keep a stable solution. The method is evaluated a few relatively simply PDEs with simple domains.