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

The writing is subpar

We appreciate the reviewer for this remark. This and other typos/errors were corrected.

The authors should compare their method with DEQ.

In our work, we discuss implicit constraints for optimization problems. While DEQ is based on special activation maps (rather than recursive functions in conventional deep learning). We don’t see a relation between these two notions.

The results in Section 3.4 is very poor and in no experiments the authors compare their method with other methods.

We redid our experiments in Section 3.4 with a new cost function, which ultimately improved inversion results. Please review this section in the updated manuscript. See also Official Comment #2 above.

First and foremost, the paper feels incomplete

We appreciate the reviewer for these remarks. Our approach is self-supervised, thus massive datasets are not needed; Training could be performed with any optimization method, either stochastic or deterministic. The work "Learning to Solve PDE-constrained Inverse Problems with Graph Networks" is not feasible for comparison since it involves solving the wave equation. For solving the wave equation explicit schemes are totally sufficient. Incorporation of explicit schemes into the AD framework was earlier discussed by Zhu el, 2021. In contrast, our work is focused on incorporation of implicit schemes. We thus did not include the wave equation in our experiments.

We added a Related Works Section and discussed the originality of our approach there.

Also See Official Comment #2 above for a comparison versus published works.

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

There is no limitation on boundary condition types since we follow the traditional FVM. For example, experiments in Section 3.4 Inverse Source Problem involve both Dirichlet and Neumann boundary conditions.

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

FVM provides the solution within the whole modelling domain. However, in many applications we have data access only at one, two or several points of the domain (e.g. the humans typically measure temperature at one point; earth subsurface properties are known only where a well was drilled). Operator R simply keeps in a given vector several entries (wherever sensors are located) while removing the others. In Figs 2a), 4a), 5a) these points are indicated as black diamonds.

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

A significant portion of inverse problems are ill-posed thus admitting an infinite number of solutions. A stabilizer S(theta) is intended to regularize an ill-posed problem, making it well-posed. The choice of S(theta) depends not only on equation type, but also on physical properties of the unknown variable as well as the measurement setup. For example, the reader may notice that different stabilizers were used in (15) and (17).

In equation (8), we need to find A^{-1} in the end. Isn't the cost for this large? Eq (8) gives only an expression for $u$. In our implementation, we solve $A u=b$ to find $u$ with a sparse direct method, thus manipulations with $A^{-1}$ are avoided.

Zhu et al, A general approach to seismic inver-sion with automatic differentiation. Computers and Geosciences, 151:104751, June 2021. ISSN 0098-3004. doi: 10.1016/j.cageo.2021.104751. URL http://dx.doi.org/10.1016/j.cageo.2021.104751

the paper would benefit from stronger experimental or theoretical justification

We extended the Into and Experimental sections and added Related Works Section

limitations of automatic differentiation in JAX.

The JAX AD solver is based on the QR matrix factorization. It is suitable for small and moderate size problems, but for larger problems it will be too slow and memory-consuming.
In our work, we show how an arbitrary solver could be integrated into the automatically differentiationable framework.

validation of conservation properties

We added the respective experiment to the manuscript, Appendix A.

Equation (12) appears to involve matrix inversion

Eq. (12) should be interpreted as solution of a linear, $A \cdot \nabla_b L = \nabla_u L,$ with a known right-hand side, $\nabla_u L$. The computational cost of this step is equal to the complexity of solving a sparse linear system.

solving the implicit equation through optimization techniques

Matrices arising from discretization of elliptic PDEs, like (5), are highly ill-conditioned. Thus incorporation of such a matrix intro optimization problems would lead to very slow convergence of the minimization algorithms (both stochastic and deterministic).

various boundary conditions

Experiments in Section 3.4 Inverse Source Problem involve both Dirichlet and Neumann boundary conditions.

provide a detailed comparison with explict Euler

We added an experiment on this topic. See Table 1 in the revised manuscript.

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.

The initial loss might be scaled weights w_i, Eq (15), so that it is 1. So far we followed the standard ML guideline to stop optimization wherever the loss is sufficiently flat. But a more precise criteria is in fact needed.

coefficient may not be unique in this setting. Could the authors clarify whether this issue arises from the problem or the numerical method?

Both problems in Fig 4 and 5 are highly under-determined. For example in Fig 4, we estimated the values of 900 unknowns based on 28. This implies non-uniqueness of the solution of the inverse problem.

Additional experimental results with fewer grid points would strengthen the paper.

We have not managed to perform such an experiment with the given time frames. But we will evidently address this topic in the future.

The novelty of the work is limited

We appreciate this remark. We added a Related Works Section and discussed the originality of our approach there. Backpropagation and the adjoint state method are known to be equivalent, see e.g. Backpropagation: Theory, Architectures, and Applications, Psychology Press, 2013.

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.

We added a comparison, see Table 1. See Official Comment #2 above for a comparison versus published works.

The presentation is not clear...

We have updated the plots for better readabilty.

What would be the limitation of the method?

Resolution of some inverse problems is quite low (e.g. due to a measurement setup). It is thus not feasible to represent the unknown variable on a near-uniform grid (Figs 4b and 5b). In our future work, we will investigate leanable grid coarsening within inverse problems.

Currently we are using a SciPy direct solver for solution Eqs (7) and (12). However, large-scale modelling and inversion would require a more economical iterative solver, which will be implemented in the future.

What would be the potential benefit of using machine learning for linear PDE over classical methods?

This is a quite general question actually. The particular problem that we addressed in our work is the incorporation of implicit constraints into the AD framework. Deriving jacobians for classical numerical methods is a colossal and prone-to-bugs task. Here it is avoided. Specifically, We designed a fully differentiable implicit solver for a linear PDE that can be used within larger pipelines, e.g. solving inverse problems with gradient methods.

Dear Reviewers,

As suggested, we performed a comparison with a baseline method.

Continuing Section 3.4, Inverse Source Problem, we compared the results of our inversion method with those obtained using physically-informed neural networks (PINNs) [1], implemented via the Deepxde package [2]. In PINNs, both the unknown solution and the right-hand side are parameterized by a neural network. For our example, we employed a fully-connected neural network with two inputs, three dense layers (150 nodes each), two outputs, and a sine activation function. The same setup described in Section 3.4 was analyzed for consistency. Both minimizations were performed using the Adam optimizer.

The figure below illustrates the results. The PINN successfully captured the primary features of the right-hand side (e); however, the source and sink were inaccurately located near the upper boundary of the domain (measurement line). Our approach was somewhat more accurate in recovery of the source and sink (b): there is a displacement above their true locations. Moreover, our method more precisely captured the magnitudes of the source and sink, than the PINN.

Lastly, we note that the convergence of the optimizer in our method was significantly faster, as shown in (d) and (f). We attribute this difference to the highly nonlinear nature of the loss functions in PINNs, whereas our approach involves a quadratic loss function (eq. 17), leading to more efficient optimization.

Figure

[1] Raissi et al, Physics-informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations. Journal of Computational Physics, 2019.

[2] Lu et al, A deep learning library for solving differential equations , SIAM Review, 2021.

Dear Reviewers,

We appreciate your constructive feedback. We would like to emphesize changes that we made in the revised manuscript.

• Added Related Work Secttion
• We added an experiment on comparison of implicit and explicit time-stepping. See Table 1 in the revised manuscript
• We added an experiment illustrating energy conservation, see Appendix A.
• We have updated the plots for better readability.

We also made a comparison versus the PINN (see Official Comment #2), which we plan to include in the camera-ready version. We hope all the reviewers will consider raising their scores further, taking into account the extensive improvements made to the paper and the additional experiments included.

Best regards,

The Authors