Mechanistic PDE Networks for Discovery of Governing Equations

We appreciate the reviewer’s comments.

Main Contribution and Paper Structure

We would like to strongly emphasize that we do not claim to replace or improve existing PDE solvers. Our contribution is to enhance the MechNN model (Pervez et al, ICML 24) that works with ODEs to support PDE representations, together with an efficient way of solving the representations and applying the model for PDE discovery. The PDE solver we develop is specialized for this architecture, where the PDE terms are produced by a (non-smooth) neural network output and solves a relaxation. The sparse solver is intended to make MechNN feasible for PDEs with multiple dimensions.

Furthermore, we make no claims of either speed or memory improvement over the baseline sparse regression methods which are not neural networks and require much fewer resources. We do claim that our model is more expressive (handles complex expressions) and can handle more complex data as shown in the discovery experiments.

We are open to restructuring the paper for clarity as long as it correctly emphasizes our main contributions and will certainly consider the reviewer’s suggestions.

Plots look similar

We disagree with the reviewer’s assessment of the plots. We claim significant improvement on the more complex reaction diffusion system (Figures 4,5 right) over the baselines which entirely fail for this dataset.

We note that the TPR and max coefficient error plots should be considered together, since the correct terms have to be discovered for the coefficient error to be meaningful. The baselines fail in all cases to discover the correct terms (with low error) including the case with no noise. For our model both TPR and coefficient error show gradual degradation with noise, with near perfect recovery in low noise settings even though we do not fine-tune after thresholding.

On the easier reaction diffusion system the WeakSINDy baseline performs well (although PDEFIND fails with noise) and our method matches its performance, with only some decrease in coefficient accuracy at the 80% noise level.

Evaluation/Simulation

The metrics of TPR and max coefficient error are chosen from the literature. In our view they provide a more quantitative and objective basis for comparison of discovery methods. Raw equations allow limited insight, especially when exact discovery is not possible in noisy settings. However, we will consider adding examples of discovered equations in the updated draft.

Similarly, simulating a discovered equation does not serve as an accurate determinant of how well an equation’s form has been learned, especially in noisy settings. Nevertheless we link a video example of simulation for the harder reaction diffusion example simulating the ground truth and discovered equation on clean data (uploaded here: https://filebin.net/wqv4t3oki732nuq0). We will consider adding more examples in the next draft.

Burger’s equation

With the inviscid Burger’s equation there is only one term to be discovered with the coefficient of 1. However, since we do not finetune after thresholding there can be occasional variation in the discovered coefficients. The learned coefficient in the noise-free case for this experiment is 0.957 which gives a relative error (|1-0.957|/|0.957|) of 0.04. For the noisy case the coefficient is 1.01 with a relative error of 0.001. Running the experiment with fine-tuning after thresholding gives a small relative error (~0.001) for the first case as well.

Recent Methods

Due to lack of space, please see our response to reviewer Ct17 for details of the comparison with PDE-LEARN.

Hyperparameters

We clarify that the baselines (PDEFIND, WeakSINDy) are sparse regression methods which do not employ any neural networks and therefore the learning rates etc, do not apply to the baselines. For our method, for simplicity, we fixed the learning rate across all our experiments and do not use any scheduling.

Retraining

Yes, the models have to be retrained for new data. In the discovery setting we consider, the PDE parameters to be discovered are global parameters and are not functions of the data. This is also true for the baselines. This setting does not allow direct evaluation of parameters on a new dataset. However, developing a setting where this is sensible is an interesting question for further work.

Threshold

We chose the threshold per PDE qualitatively, for simplicity. The choice of threshold can affect small parameters which is also the reason why the TPR in the Navier-Stokes example is less than 1.

Analytical solutions

Most PDEs don’t have analytical solutions, but perhaps the question is with regard to learning the expression structure? If so, then yes, we think that symbolic regression methods could be combined with this method to learn expressions together with parameters and we are considering approaches for this.

We would like to thank the reviewer again for the review.

We thank the reviewer for the comments.

The reviewer’s point regarding use of the term ‘discovery’ is well received. However, the term is now endemic in the literature and we chose to employ the same term.

We would like to appreciate the reviewer’s recognition of the effort that went into building the solver. We will certainly release the source code for our models with the camera ready version of the paper.

Multigrid

It is quite correct that multigrid solvers are also used for higher resolutions. Our use case is different in a few respects from the usual use of multigrid.

• One difference is that we solve for the saddle point systems in equations 10,11 and we have extra constraints for ensuring smoothness which are usually not present in solvers.
• Secondly, we work in a learning context with a nested iterative loop. The inner loop corresponds to the multgrid iteration and the outer loop is the learning loop that can run for hundreds of epochs.
• Similarly we work in a batch setting where in each inner iteration we solve a minibatch of PDEs, all with their separate multgrid steps.
• Furthermore the backward pass also uses the batched multigrid solver which doubles the memory use.
• It is known that multigrid performance deteriorates with discontinuous coefficients. The data produced by neural networks over the grid (especially early in training) can be non-smooth which can deteriorate multigrid performance.
• Our current implementation is in Pytorch which is not ideal for iterative processing. A native CUDA implementation can potentially produce significant computational improvements, but requires further work.

Very high multigrid resolutions would have an adverse impact on the walltime, making the method very slow. Nevertheless, different from standard multigrid approaches, since we also have a learning component which can infer the underlying PDE given only a (batched) portion of the full resolution data. This implies that the solver is not required to solve for the full resolution data in each step. Rather it only has to solve for the given smaller resolution patch of data making it more efficient.

Encoder NN

The neural network transform (line 267) allows parameter capacity to the model to make it flexible. Without the transformation the only learnable parameters would be the PDE parameters, parameterizing the expression, which are few in number. This can lead to a brittle model making it harder to recover from bad initialization or noise. We will consider adding an ablation in the next draft.

$x_1 x_1$ should be $x_1 x_2$. Yes, that is a typo in this case.

We would like to thank the reviewer again for the review.

We thank the reviewer for the comments. We attempt to address the concerns raised below.

Non-autonomous systems

Yes, in this paper we focus on autonomous systems. However, non-autonomous systems can be similarly handled since the method allows spatiotemporally varying terms. Any known external forces can be represented as time-space dependent terms similarly to what is achieved with basis functions. Any unknown forces that cannot be simply represented can be represented by a neural network. This is possible because the method allows arbitrary differentiable expressions in PDEs.

Identifiability

Whether the method finds the unique PDE requires further investigation into identifiability of the PDE learning problem. Identifiability is something we do not deal with in this paper and assume that the problems are identifiable and that there is a unique governing equation.

We have not considered stochastic or fractional differential equations thus far and it would be interesting to consider whether such types of equations can be modeled in this way.

PINN style methods

The PDE-LEARN approach extends the PINN approach to PDE discovery. It reduces a PDE residual, constructed using basis functions, at random collocation points together with approximate L0 norm regularization. One difference of this approach with MechNN is that MechNN has a stronger physics informed bias due to the explicit modeling of PDEs as constraints.

Unfortunately, the PDE-LEARN release does not have support for coupled equations, so it cannot be used with our 2D experiments in its current form. The main equations that we focused on were the 2D reaction diffusion and Navier-Stokes PDEs which are coupled equations.

Nevertheless, we trained PDE-LEARN on our 1D Burger’s equation dataset with a 1/cosh(x) initial condition. Despite repeated attempts with multiple settings of hyperparameters, we found it to always over or under prune the equation (even on clean data), performing worse than the SINDy baselines from our paper. The maximum TPR we obtain for this experiment is 0.5, compared with a TPR of 1 for the SINDy baselines and MechNN.

In particular one advantage of our method is that the modeled PDEs can contain arbitrary differentiable expressions and are not limited to linear combinations of fixed basis functions. Many other methods including PDE-LEARN model PDEs as linear combinations of fixed basis functions. A concrete example is the porous medium equation for which the equation cannot be modeled in this way, whereas we show that we can recover the true parameters.

Limitations

The following are some limitations of the method.

• The method is slower than the baseline sparse regression methods.
• The method is currently limited to Cartesian grids.
• Although the method can represent arbitrary differentiable expressions, there remains the need for theory that dictates the conditions under which complex expressions can be exactly identified.
• There is a tradeoff between accuracy and speed in the multigrid solver which needs to be tuned.
• Application to very high dimensional data is not feasible in the current form.

We thank the reviewer again for the review and hope that we were able to meet any concerns.

We appreciate the detailed review.

Clarification of our contribution.

We wish to clarify the nature of our contribution. We do not claim to replace or improve existing PDE solvers. Our contribution is to enhance the MechNN model (Pervez et al, ICML 24) that works with ODEs to support PDE representations, together with an efficient way of solving the representations and applying the model for PDE discovery. The PDE solver we develop is specialized for this architecture, where the PDE terms are produced by a (non-smooth) neural network and solves a relaxation. The sparse solver is intended to make the MechNNs feasible for PDEs with multiple dimensions. That said, we do not discount the possibility that the method could be useful for other applications.

Furthermore, we make no claims of either speed or memory improvement over the baseline sparse regression methods which require much fewer resources. We do claim that our model is more expressive (handles complex expressions) and can handle more complex data as shown in the discovery experiments.

Missing Evidence and Failure.

We provide quantitative evidence of discovery performance using the TPR and Error metrics from the literature. A TPR of 1 means exact discovery of terms. This is shown in Figures 4,5 and 6 for our method and the baselines. The plots also show partial failure to discovery when noise is added. For instance with 80% noise (Figure 4) we only see a TPR of 0.6 which indicates extra or missing terms are present. From Figure 6 we see that the Navier-Stokes equation is also not exactly discovered (TPR ~0.8). We present quantitative metrics since they allow more objective comparison than qualitative comparison between equations. We will include examples of discovered equations.

Solver Eval

We have performed further evaluation of the multigrid linear solver including time and memory usage, convergence of relative error with the number of V-cycles, GS smoothing iterations, and FGMRES iterations for the 2D Laplace equation. See https://filebin.net/wqv4t3oki732nuq0 for plots.

GPU Multigrid

GPU multigrid methods are frequently used in PDE solvers. However, our solver is specialized to work with (possibly non-smooth) neural network outputs over the grid that parameterize the PDEs. The non-smoothness is handled by adding extra smoothing constraints and solving relaxations as saddle-point problems. We are not aware of available solutions for the special case that we consider.

We plan to release the code for our experiments.

Comments

By grid size we mean increasing the grid resolution for the same problem and running until convergence. We will make the correction.

The line about derivatives is with reference to the SINDy baselines which estimate numerical derivatives directly on data which makes them susceptible to noise. There is a typo and the word ‘computed’ is missing.

The statement about the parameter $m$ was meant to indicate that it is not a fixed value which implies that fixed basis functions cannot be used to represent the PDE.

We will make the corrections in the next draft.

Questions

The V-cycle algorithm itself is fairly standard, however, it is applied in a non-standard setting to solve saddle-point problems in eq 10 and 11. For these problems the standalone V-cycle has subpar performance in our setting. However, we were able to obtain better results using the V-cycle as a preconditioner for FGMRES.

Yes, we also solve the backward pass with the V-cycle (with FGMRES).

We will include more details about the solver implementation in the next draft. The code is in Pytorch together with CuPy for features not available in Pytorch (such as sparse triangular solves). We did not write CUDA kernels, however, that is an option for further improvement. For FGMRES we start with the CPU implementation of GMRES from SciPy. We adapted the GMRES algorithm to FGMRES using the algorithm described in Saad (2003). We ported the code to Pytorch adding batching and GPU support and V-cycle preconditioning, written from scratch.

For simplicity we only threshold once at the end of training (unlike SINDy). The threshold is a hyperparameter per dataset (fixed across noise). We used 0.06, 0.3 and 0.0002 for the easy and hard reaction diffusion, and Navier-Stokes equations. Fine-tuning after thresholding can be used to obtain more accurate coefficients.

The parameterization with a ResNet refers to computing \tilde{u} (line 267 or Figure 1). We parameterize the PDE terms (Figure 1) by a neural network (using \tilde{u}) instead of feeding data (u_data) (which may be noisy) directly. This also helps with training convergence.

Line 411 discusses discovery in the presence of noise. Results are shown in Figure 4 (left). WeakSINDy and MechNN have similar performance with at least TPR 0.6 whereas PDEFIND fails entirely.

We thank the reviewer again.