Differentiable Implicit Solver on Graph Neural Networks for Forward and Inverse Problems

Authors: Nikolay Yavich, Alexander Ryabov, Evgeny Burnaev, Vladimir Vanovskiy
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
Keywords: graph neural networks, differentiable solvers, implicit schemes, numerical modelling, inverse problems

Abstract:
Partial differential equations (PDEs) on unstructured grids can be solved using message passing on a graph neural network (GNN). Implicit time-stepping schemes are often favored, especially for parabolic PDEs, due to their stability properties. In this work, we develop a fully differentiable implicit solver for unstructured grids. We evaluate its performance across four key tasks: a) forward modeling of stiff evolutionary and static problems; b) the inverse problem of estimating equation coefficients; c) the inverse problem of estimating the right-hand side; and d) graph coarsening to accelerate forward modeling. The increased stability and differentiability of our solver enable excellent results in reducing the complexity of forward modeling and efficiently solving related inverse problems. This makes it a promising tool for geoscience and other physics-based applications.

Decision: reject
Reviewer scores: 3, 1, 3, 1

Reviews and discussion:
<thread>
creation date:
1730576585111

thread_type:
official_review

<thread_object>
creator: reviewer_2k7b

<content>
summary:
This paper explores the use of graph neural networks for solving forward and inverse problems and particularly focuses on the incorporation of implicit solver. However, the writing is subpar and the procedures and advantages are not well explained. The experiments are also lack comparison with other methods.

strengths:
The question is interesting and combining GNN with finite element method seems natural.

weaknesses:
1. The writing is subpar. There are many typos and grammatical errors. For example, "Compute $\nabla_bL$  with (12), whats is equivalent so the solution of a single linear system." should be "Compute $\nabla_bL$ with (12), which is equivalent to solving a single linear system."
2. One main focus of this paper is the incorporation of implicit solver. However, using an iterative solver and in a deep learning setting is well-studied in the Deep Equilibrium Models (DEQ) literature. The authors should compare their method with DEQ.
3. The experiments are not very convincing. The results in Section 3.4 is very poor and in no experiments the authors compare their method with other methods.

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

</thread>
<thread>
creation date:
1730686229000

thread_type:
official_review

<thread_object>
creator: reviewer_qrey

<content>
summary:
This paper proposes an integrated approach for solving forward and inverse problems by creating a new pipeline that combines Graph Neural Networks (GNNs) with Finite Volume Methods (FVM) to enable automatic differentiation with implicit-time stepping.

strengths:
Figure 1 effectively illustrates the overall pipeline, demonstrating experimental results that apply the combination of GNN and FVM to both forward and inverse problems.

weaknesses:
First and foremost, the paper feels incomplete. The biggest concern is the lack of discussion about other approaches that use GNNs or integrate FVM with deep learning to solve PDEs. A “Related Work” section should be added to explain how the proposed model differs from recent studies and highlight its novelty. Although Section 2 on theory explains the problem setup to some extent, more detailed steps and methods for training the proposed approach should be included. Section 3, the experimental part, merely lists the results for forward and inverse problems without discussing how this method compares to existing GNN- and FVM-based approaches. For instance, the study "Learning to Solve PDE-constrained Inverse Problems with Graph Networks" solves inverse problems using GNNs—how does the proposed method differ from this approach, and what advantages does it offer? Experimentally, does it outperform in solving inverse problems?

questions:
* Why should we only consider the Neumann boundary in equation (1)? Is it difficult to consider other Robin boundaries?

* I don't understand the role of R in equation (4). What does it mean as a measurement operator?

* Does S(theta) change depending on the equation of the PDE to be solved? Can you explain this further?

* In equation (8), we need to find A^{-1} in the end. Isn't the cost for this large?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730720190620

thread_type:
official_review

<thread_object>
creator: reviewer_1yfg

<content>
summary:
The work considers mesh coarsening, forward, and inverse problems and investigates the implicit solver. In the numerical experiments, the authors evaluated the performance of the method regarding each problem.

strengths:
* The work tries to build a framework that works with mesh coarsening, forward, and inverse problems.

weaknesses:
* The novelty of the work is limited. Incorporating FVM into GNN is not new and considered in, e.g., [Jessica et al. ICML 2024 https://arxiv.org/abs/2311.14464 ] and [Horie et al. ICML 2024 https://arxiv.org/abs/2405.16183v1 ]. The construction of gradients presented in Section 2.3 seems strongly related to the adjoint method, which is a standard way to deal with inverse problems. The implicit method for GNN is considered in the area of implicit GNNs, e.g., [Gu et al. NeurIPS 2020 https://arxiv.org/abs/2009.06211 ]. The authors state that these are their novelty, but there is existing work for each. The authors should cite these works and clarify the added novelty from the authors.
* The evaluation is weak. There is only one baseline for the experiment in Section 3.2 and nothing for the ones in Section 3.3 and 3.4. With the current form, the reviewer cannot asses the effectiveness and superiority of the model.
* The presentation is not clear. The figure may miss the labels (a), (b), and so on for Figures 2, 3, and 4. It is not clear what is "data 1", "fitting 1", "data 2", and "fitting 2" in Figures 2 and 3.

questions:
* What would be the limitation of the method?
* What would be the potential benefit of using machine learning for linear PDE over classical methods?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730721454224

thread_type:
official_review

<thread_object>
creator: reviewer_4kqz

<content>
summary:
This paper introduces a novel framework that combines graph neural networks with the finite volume method, to address implicit schemes. (There was a challenge that differential equation solvers typically avoid due to the additional computation complexity of handling implicit equations.)

strengths:
By employing an implicit scheme with optimized gradient computation, the proposed method reduces the required number of time steps. 
They present a differentiable framework for both forward and inverse methods, enabling a learnable numerical approach based on discrete time steps. 
Additionally, the paper explores applications in inverse problems, often employing irregular unstructured grids as used in practical scenarios.

weaknesses:
While the underlying idea is promising, the paper would benefit from stronger experimental or theoretical justification for the proposed methodology. Additional clarity and motivation for the approach would enhance the paper’s impact.

questions:
1. In line 74, the authors discuss the limitations of automatic differentiation in JAX. Further elaboration on this limitation would improve the motivation for this approach.

2. Graph neural networks are frequently used to manage unstructured grid points. Since the paper emphasizes integration with the finite volume method with its local conservation property in line 58, it would be beneficial to include experiments validating these conservation properties. 

3. Equation (12) appears to involve matrix inversion on the right-hand side of the proposed gradient formulation. Could the authors address whether this matrix inversion contributes to computational costs, comparable to previous methods?

4. In Equations (12) and (13), gradient computations are proposed. An alternative approach might involve solving the implicit equation through optimization techniques commonly used in deep learning, such as constructing a minimization problem for Equation (5) combined with a data loss function, potentially avoiding matrix inversion. Could the authors discuss this approach?

5. While the finite volume method can accommodate various boundary conditions, the paper considers only Neumann boundary conditions. Is there a specific reason for this choice?

6. In line 253, the authors claim lower computational costs for their method than the explicit Euler scheme, which requires smaller time steps. Could the authors provide a detailed comparison of the computational cost per each time step to support this claim?

7. In Figure 2, the initial scale of the loss is relatively high, making it difficult to assess whether the loss converges to zero after 160 epochs. Given that a coarser grid, intended to reduce computation, may negatively impact estimation accuracy even if the loss function converges to zero, a guideline for determining sufficient loss minimization would be beneficial.

8. In Figures 3 and 4, the recovered permeability only captures general trends rather than precise values. However, the proposed method in (c) accurately approximates the true data distribution, which suggests that the problem may be inherently ill-posed, where the coefficient may not be unique in this setting. Could the authors clarify whether this issue arises from the problem or the numerical method?

9. All experiments utilize a large number of data points, which may facilitate finding a solution. Additional experimental results with fewer grid points would strengthen the paper.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732837130610

thread_type:
official_comment

<thread_object>
creator: authors

<content>
comment:
Dear Reviewers,

We appreciate your remarks. 
We should emphasize that our work focuses on incorporation of implicit constraints into the AD framework. 
The main difference versus several main-stream papers is that
* our approach is self-supervised, thus massive datasets are not needed;
* our approach does not produce a surrogate solver, but rather follows the FVM discretization respecting energy conservation.
* In particular, our approach is quite close to the traditional method of solution of inverse problems, i.e. Tikhonov functional minimization. However, we employed AD to avoid derivation of jacobians and GNN for unstructured data representation. The challenge was to implement AD efficiently and flexibly whenever implicit constraints are given and described by large sparse matrices

We revised the manuscript to address the issues and added new experiments.
The new text is in red for your convenience. 
 
If there are any further questions or points that need a discussion, we will be happy to address them. 
 
 Best regards,

 The Authors
</content>
</thread_object>

</thread>


Metareview:
metareview:
(a) The high level idea pursued in this paper is to combine finite volume methods, graph neural networks (GNNs), and implicit time-stepping to solve partial differential equations on unstructured grids. (b) The idea is intuitive but (c) is supported by a limited set of experiments on toy datasets. (d) Scientific machine learning is now a rather mature sub-field, and in my opinion any new techniques, even if intuitive,  should be supported with solid evaluations on a diverse set of PDE systems, forcing functions, boundary conditions, etc. This paper falls (significantly) short in this regard. The original version of this paper also didn't have any discussion about connections with related work.