Learning a Neural Solver for Parametric PDE to Enhance Physics-Informed Methods

## Abstract

Abstract Physics-informed deep learning often faces optimization challenges due to the complexity of solving partial differential equations (PDEs), which involve exploring large solution spaces, require numerous iterations, and can lead to unstable training. These challenges arise particularly from the ill-conditioning of the optimization problem, caused by the differential terms in the loss function. To address these issues, we propose learning a solver, i.e., solving PDEs using a physics-informed iterative algorithm trained on data. Our method learns to condition a gradient descent algorithm that automatically adapts to each PDE instance, significantly accelerating and stabilizing the optimization process and enabling faster convergence of physics-aware models. Furthermore, while traditional physics-informed methods solve for a single PDE instance, our approach addresses parametric PDEs. Specifically, our method integrates the physical loss gradient with the PDE parameters to solve over a distribution of PDE parameters, including coefficients, initial conditions, or boundary conditions. We demonstrate the effectiveness of our method through empirical experiments on multiple datasets, comparing training and test-time optimization performance 1 1 1 Code will be released soon at https://github.com/2ailesB/neural-solver .

## 1 Introduction

Partial Differential Equations (PDEs) are ubiquitous as mathematical models of dynamical phenomena in science and engineering. Solving PDEs is of crucial interest to researchers and engineers, leading to a huge literature on this subject (Evans, ; Salsa, ). Traditional approaches to solving PDEs such as finite difference, finite element analysis or spectral methods (Zienkiewicz et al., ; LeVeque, ) often come with stability and convergence guarantees, but suffer from a high computational cost. Improving numerical PDE solvers through faster and more accurate algorithms remains an active research topic (Zienkiewicz et al., ).

PDE solvers usually rely on discretization and/or linearization of the problem through various techniques, to simplify the computations. Iterative methods such as Jacobi, Gauss-Seidel, Conjugate Gradient, and Krylov subspace methods can then be used to solve the resulting systems. Unfortunately, many PDEs have an ill-conditioned nature and these iterative processes can demand extensive computational resources. Preconditioning techniques are often essential to mitigate this, though they require precise customization to the specific PDE problem, making the development of effective solvers a significant research endeavor in itself. Yet, the computational demands, time, and expertise required to develop these algorithms sometimes make them impractical or sub-optimal for specific classes of problems. Instead of relying on hand-designed algorithms, researchers have investigated, as an alternative, the use of machine learning for training iterative PDE solvers Hsieh et al. (); Li et al. (); Rudikov et al. (); Kopaničáková et al. (). These approaches usually parallel the classical numerical methods by solving a linear system resulting from the discretization of a PDE, for example, using finite differences or finite elements. A preconditioner is learned from data by optimizing a residual loss computed w.r.t. a ground truth solution obtained with a PDE solver. This preconditioner is used on top of a baseline iterative solver and aims at accelerating its convergence. Examples of baselines solvers are the conjugate gradient (Li et al., ; Rudikov et al., ) or the Jacobi method Hsieh et al. ().

Another recent research direction investigates the use of neural networks for building surrogate models in order to accelerate the computations traditionally handled by numerical techniques. These methods fall into two main categories: supervised and unsupervised. The supervised methodology consists of first solving the PDE using numerical methods to generate input and target data and then regressing the solution using neural networks in the hope that this surrogate could solve new instances of the PDE. Many models, such as Neural Operators, lie within this class (Li et al., ; Raonić et al., ; Bartolucci et al., ) and focus on learning the solution operator directly through a single neural network pass.
Unsupervised approaches, involve considering a neural network as a solution of the PDE. The neural network parameters are found by minimizing the PDE residual with gradient descent. Methods such as Physics-Informed Neural Networks (PINNs) (Raissi et al., ), or DeepRitz (E & Yu, ) fall under this category. This family of methods is attractive as it does not rely on any form of data, but only on information from the PDE residual. However, they exhibit severe difficulties for training (Krishnapriyan et al., ; Ryck et al., ), often requiring many optimization steps and sophisticated training schemes (Krishnapriyan et al., ; Rathore et al., ). The ill-conditioned nature of PDEs residual loss appears again in this context, making standard optimizers such as Adam inappropriate. A more detailed review of the existing literature is described in appendix A.

In this work, we consider having access to the PDE as in unsupervised approaches and also to some data for training our neural solver. Our objective is to solve the optimization issues mentioned above by learning an iterative algorithm that solves the PDE from its residual defined as in the PINNs framework (see fig. 1). This neural solver is trained from data, either simulations or observations. Different from the classical ML training problem which aims at learning the parameters of a statistical model from examples, the problem we handle is learning to learn Andrychowicz et al. (), i.e. learning an iterative algorithm that will allow us to solve a learning problem. When vanilla PINNs handle a single PDE instance, requiring retraining for each new instance, we consider the more complex setting of solving parametric PDEs, the parameters may include boundary or initial conditions, forcing terms, and PDE coefficients. Each specific instance of the PDE, sampled from the PDE parameter distribution, will then be considered as a training example. The objective is then to learn a solver from a sample of the parametric PDE distribution in order to accelerate inference on new instances of the equation. With respect to unsupervised approaches, our model implements an improved optimization algorithm, tailored to the parametric PDE problem at hand instead of using a hand-defined method such as stochastic gradient descent (SGD) or Adam. As will be seen in the experimental section, this idea will prove extremely effective for the ill-posed problem of optimizing PINNs objective, allowing convergence in a few steps. In the proposed methodology, the neural solver will make use of gradient information computed by a baseline gradient method in order to accelerate its convergence. In our instantiation, we will use SGD as our baseline algorithm, but the method could be easily extended to other baselines. Our model deviates from the traditional preconditioning methods by directly optimizing the non-linear PDE residual loss (Raissi et al., ) without going through the discretization steps.

Figure: Figure 1: Optimization scheme of a physics-informed method with our framework.
Refer to caption: x1.png

Our contribution includes :

- •
Setting an optimization framework for learning to iteratively solve parametric PDEs from physics-informed loss functions. We develop an instantiation of this idea using an SGD baseline formulation. We detail the different components of the framework as well as training and inference procedures.
- •
Evaluating this method on challenging PDEs for physics-informed method, including failure cases of classical PINNs and showing that it solves the associated optimization issues and accelerates the convergence.
- •
Extending the comparison to several parametric PDEs with varying parameters from $1$d static to $2$d+time problems. We perform a comparison with baselines demonstrating a significant acceleration of the convergence w.r.t. baselines.

## 2 Motivation

Let us first motivate our objective with a simple example. Physics-informed neural networks (PINNs) are a promising tool for optimizing neural networks in an unsupervised way to solve partial differential equations (PDEs). However, these methods are notoriously difficult to train (Krishnapriyan et al., ; Ryck et al., ). As an illustrative example of this challenge, let us solve the zero-boundary Poisson equation in $1$d on $\Omega=[-\pi,\pi]$.

##### Poisson equation, 1d.

The solution is given by $u(x)=\sin(kx)$.

$$ $\displaystyle u^{\prime\prime}(x)=-k^{2}\sin(kx),$ $\displaystyle u(-\pi)=0,\,u(\pi)=0.$ (1) $$

Physics-informed machine learning relies on an *ansatz space* of *parametric functions*, e.g. neural networks $u_{\Theta}:\Omega\mapsto\mathbb{R}$, minimizing the following loss in order to satisfy the constraints in equation 1:

$$ $\mathcal{L}_{\textnormal{PDE}}=\mathcal{L}_{{}_{\textnormal{Res}}}+\lambda \mathcal{L}_{\textnormal{BC}},\quad\mathcal{L}_{{}_{\textnormal{Res}}}=\int_{ \Omega}\left|u_{\Theta}^{\prime\prime}(x)-f(x)\right|^{2}dx,$ $$

$$ $\mathcal{L}_{\textnormal{BC}}=\frac{1}{2}\left[u_{\Theta}(-\pi)^{2}+u_{\Theta} (\pi)^{2}\right].$ $$

As a simple example, consider the parametrization given by considering a linear combination of Fourier features widely used (Tancik et al., )(^2^22Note that even though the ansatz is linear in $\Theta$, it is not linear in $x$.), $u_{\Theta}(x)=\sum_{k=-K}^{K}\theta_{k}\phi_{k}(x)$, with $\phi_{0}(x)=\tfrac{1}{\sqrt{2\pi}}$, $\phi_{-k}(x)=\tfrac{1}{\sqrt{\pi}}\cos(kx)$ and $\phi_{k}(x)=\tfrac{1}{\sqrt{\pi}}\sin(kx)$ for $1\leq k\leq K$.

This simple but informative example yields a tractable gradient descent algorithm, as the associated updates are linear in the parameters, governed by a matrix $A$ and constant $b$:

$$ $\displaystyle\Theta_{l+1}$ $\displaystyle=\Theta_{l}-\eta\nabla\mathcal{L}_{\textnormal{PDE}}(\Theta_{l})$ $\displaystyle=(I-\eta A)\Theta_{l}+b$ (2) $$

with $A$ whose condition number is $\kappa(A):={\lambda_{\max}(A)}/{\lambda_{\min}(A)}\geq K^{4}$:

$$ $\displaystyle A$ $\displaystyle=\begin{bmatrix}0^{4}&0&\cdots&0\\ 0&1^{4}&\cdots&0\\ \vdots&\vdots&\ddots&\vdots\\ 0&0&\cdots&K^{4}\end{bmatrix}+\lambda\begin{bmatrix}\phi_{1}(\pi)\\ \phi_{2}(\pi)\\ \vdots\\ \phi_{K}(\pi)\end{bmatrix}\begin{bmatrix}\phi_{1}(\pi)&\phi_{2}(\pi)&\cdots& \phi_{K}(\pi)\end{bmatrix}.$ (3) $$

This implies that the condition number of $A$ increases extremely rapidly in the ratio between the highest and lowest frequency of the network. Given that the rate of convergence to the optimum $\Theta^{\ast}=\Theta_{0}+A^{-1}b$ can be bounded as

$$ $\displaystyle\|\Theta_{l}-\Theta^{\ast}\|_{2}\leq\left(1-{c}/{\kappa(A)}\right )^{l}\|\Theta_{0}-\Theta^{\ast}\|_{2}.$ (4) $$

The number of steps $N(\varepsilon)$ required to obtain an error of size at most $\varepsilon$, i.e., $\|\Theta_{l}-\Theta^{\ast}\|_{2}\leq\varepsilon$ increases linearly in the condition number, i.e. as the fourth power of the maximal frequency $K$:

$$ $N(\varepsilon)=O\left(\kappa(A)\ln\tfrac{1}{\epsilon}\right)=O\left(K^{4}\ln \tfrac{1}{\epsilon}\right).$ (5) $$

We believe that this simple example clearly illustrates and highlights the fact that PINNs–even when considering a linear basis, and when the PDE is linear–suffer heavily from ill-conditioning: if $500$ steps are required in order to achieve a given error when $K=5$, roughly speaking, $312\>500$ steps are required for only $K=25$.

Thus, our objective in the following will be to accelerate the convergence of such systems in this context–as well as extend them to the non-linear setting. To do so, in the following section, we will learn how to transform the optimization problem in such a way that the number of gradient descent iterations is small. The resulting method can be seen as a standalone, iterative solver as it is not only applicable to different PDEs but can handle a wide range of initial/boundary conditions and parameters.

## 3 Approach

In order to optimize PDE-based losses, we propose to learn a physics-based optimizer that will fulfill two objectives: (i) allowing a fast test-time optimization given a new PDE and (ii) solving without retraining parametric PDEs, with varying PDE coefficients $\gamma$ (^3^33Note that PDE coefficients can be functions, an example is the Darcy PDE in section 4.1), forcing terms $f$, and initial/boundary conditions $g$ using the same model. We present the general framework below and propose an instantiation that leverages a linear combination of basis functions as the ansatz.

### 3.1 Problem statement

Let us consider the following family of boundary value problems parameterized by $\gamma$ with domain $\Omega$, representing both space and time, with $\mathcal{N}$ a potentially nonlinear differential operator, $\mathcal{B}$ the boundary operator, $g$ the initial/boundary conditions, and source term $f$:

$$ $\displaystyle\mathcal{N}(u;\gamma)$ $\displaystyle=f\quad\text{in }\Omega,$ (6) $\displaystyle\mathcal{B}(u)$ $\displaystyle=g\quad\text{on }\partial\Omega.$ (7) $$

Note that different PDEs can be represented in this form, amounting to changing the parameters $\gamma$. The goal here is to develop a generic algorithm that is able to solve the above problem, yielding an approximate solution $u$ given the PDE and different sets of inputs $(\gamma,f,g)$.

For training, we assume access to a dataset of $M$ problem instances, represented by the PDE parameters $(\gamma_{i},f_{i},g_{i})_{i=1}^{M}$ and to associated target solution $(u_{i})_{i=1}^{M}$ given on a $m$ point grid $(x_{j})_{j=1}^{m}$. The solutions $(u_{i})_{i=1}^{M}$ will be used to train the neural solver. At inference, for a new PDE instance, only the PDE parameters are provided and we do not access solution points $(u_{i})_{i=1}^{M}$.

### 3.2 Methodology

Physics-informed neural networks consider an ansatz $u_{\Theta}$ parametrized by some finite-dimensional $\Theta$. The parameters $\Theta$ are iteratively updated by minimizing a criterion $\mathcal{L}_{\textnormal{PDE}}$ (e.g. the PDE residual), which assesses how well the ansatz $u_{\Theta}$ meets the conditions specified in equations 6 and 7. As introduced for example for PINNs (Raissi et al., ) or Deep Galerkin method Sirignano & Spiliopoulos (), we consider $\mathcal{L}_{\textnormal{PDE}}$ to be given by the strong formulation of the residual $\mathcal{L}_{\textnormal{Res}}$, plus a boundary discrepancy term $\mathcal{L}_{\textnormal{BC}}$(^4^44Note that other formulations of the loss may also be considered in a straightforward manner.).

$$ $\mathcal{L}_{\textnormal{PDE}}=\mathcal{L}_{{}_{\textnormal{Res}}}+\lambda \mathcal{L}_{\textnormal{BC}},\quad\mathcal{L}_{{}_{\textnormal{Res}}}=\sum_{x _{j}\in\Omega}|\mathcal{N}(u_{\Theta};\gamma)(x_{j})-f(x_{j})|^{2},\quad \mathcal{L}_{\textnormal{BC}}=\sum_{x_{j}\in\partial\Omega}|\mathcal{B}(u_{ \Theta})(x_{j})-g(x_{j})|^{2}$ (8) $$

As illustrated in section 2, performing gradient descent, or alternatives such as Adam and L-BFGS on such a highly ill-conditioned loss $\mathcal{L}_{\textnormal{PDE}}$ leads to severe training difficulties (Krishnapriyan et al., ).
The key idea in our work is to improve baseline gradient descent algorithms with the neural solver. More precisely starting from a baseline gradient algorithm, SGD in our instantiation, instead of considering the classical update, we first transform the gradient using a neural network $\mathcal{F}_{\varrho}$ with parameters $\varrho$, depending on the values of the PDE parameters as well as on other inputs such as the residual gradient provided by SGD: $\nabla_{\Theta}\mathcal{L}_{\textnormal{PDE}}$. The objective is to transform, through the neural solver $\mathcal{F}_{\varrho}$, the ill-conditioned problem into a new, simpler problem that requires fewer steps in order to achieve a given error.

Once the neural solver $\mathcal{F}_{\varrho}$ has been learned, inference can be performed on any new PDE as follows (see inference algorithm 1). Starting from an initial ansatz parameter $\Theta_{0}$, it is iteratively updated by this solver. At iteration $l$, the steepest direction of the loss $\mathcal{L}_{\textnormal{PDE}}$ is first computed with autograd. Then, the gradient is transformed, in a PDE parameter dependant way, with $\mathcal{F}_{\varrho}$:

$$ $\Theta_{l+1}=\Theta_{l}-\eta\mathcal{F}_{\varrho}(\nabla_{\Theta}\mathcal{L}_{ \textnormal{PDE}}(\Theta_{l}),\;\gamma,f,g)$ (9) $$

The objective is to iteratively refine the ansatz to closely approximate the true solution after a series of $L$ iterations, ideally small for efficiency.

This approach can be seen as learning the iterates in a PDE solver to achieve a low loss, similar to the residual minimization methods in PDEs (Elman et al., ). By design, this solver is intended to be applicable to different PDEs, as well as various sources, boundary conditions, and initial conditions. This flexibility allows for a broad range of applications, making it a versatile tool in solving complex PDEs with varying characteristics.

Designed as a parametric PDE solver, $\mathcal{F}_{\varrho}$ is trained from input target data from different sets of PDE parameters, as outlined in section 3.1. Once trained, it will be used without retraining on new instances of the PDE, i.e. with new values of the PDE parameters. The underlying hypothesis is that even though the solutions may be different for different inputs and parameters, the solution methodology remains relatively consistent. This consistency is expected to enhance the algorithm’s ability to generalize across novel scenarios effectively.

### 3.3 Training of a physics-informed solver

Choice of Ansatz $u_{\Theta}$. A very common choice (Shen et al., ) is to consider a family of basis functions $\Psi(x)=\{\psi_{i}(x)\}_{i=1}^{N}$ and consider the ansatz to be given by its linear span ${u_{\Theta}(x)=\sum_{i=0}^{N}\theta_{i}\psi_{i}(x)}$. In the following, we consider this linear reconstruction, although our formulation is generic in the sense that it can also accommodate nonlinear variants. (^5^55Although in our preliminary results we have found this may further complicate training.)

As indicated in eq. 10, the solver $\mathcal{F}_{\varrho}$ will be trained from samples of the PDE parameter distribution $(\gamma,f,g)$ and from the associated samples of the solution $u$. We first describe below the inference step aiming at iteratively updating the parameters $\Theta$ of the solution function $u_{\Theta}$ while $\mathcal{F}_{\varrho}$ is hold fixed. We then describe how the ${\varrho}$ parameters of the solver are trained. Please refer to fig. 1.

Figure: Algorithm 1 Inference

Inference
The inference step is performed at fixed values of the solver parameters $\mathcal{F}_{\varrho}$. It consists, for a given instance of the PDE with parameters $(\gamma_{i},f_{i},g_{i})$, in finding the best $\Theta$ with a few steps of the solver using eq. 9. It is illustrated in fig. 1 - grey box: starting from initial parameters $\Theta_{0}$, we compute the Physical loss $\mathcal{L}_{\textnormal{PDE}}$ using the ansatz $u_{\Theta_{0}}$. The PDE derivatives, in $\mathcal{L}_{\textnormal{PDE}}$, can be computed by hand or automatic differentiation depending on the application (^6^66In our experiments, we computed the derivative by hand when possible since it fastens computations.).
Then, eq. 9 is used to update the parameters $\Theta$ for a given number of steps $L$.
The final solution is reconstructed using the linear combination, ${u_{\Theta}(x)=\sum_{i=0}^{N}\theta_{i}\psi_{i}(x)}$, introduced in section 3.2 with the computed coefficients $\Theta_{L}$. Note that inference does not make use of the sampled
target solutions $(u_{i})_{i=1}^{M}$ computed on the grid points$(x_{j})_{j=1}^{m}$.These targets are used exclusively for training the neural solver.

Figure: Algorithm 2 Training algorithm for learning to optimize physics-informed losses.

Training the neural solver
Training amounts at learning the parameters of the solver $\mathcal{F}_{\varrho}$ and is performed with a training set of PDE parameters and simulation data considered as ground truth $(\gamma_{i},f_{i},g_{i},u_{i})_{i=1}^{M}$, corresponding to PDE instances (i.e. with different parameters $\gamma$ and/or forcing terms $f$ and/or initial/boundary conditions $g$). See fig. 1 - white box. The objective is to learn a solver $\mathcal{F}_{\varrho}$ able, at inference, to converge to a target solution in a small ($2$ to $5$ in our experiments) number of steps. For that, an optimizer (Adam in our experiments) is used to update the $\mathcal{F}_{\varrho}$ parameters. The training algorithm makes use of the data associated to the different PDE instances by sampling PDEs in batches and running algorithm 1 on several PDE instances. For each PDE instance, one starts from an initial parameter value $\Theta_{0}$ and then performs two optimisation steps (see algorithm 2): (i) one consists in solving in the ansatz parameters $\Theta$ using the neural solver using algorithm 1, leading to $\Theta_{L}$; the second one is the optimization of the solver parameters $\varrho$.
We train the outputs directly to match the associated ground truth $(u_{i}(x_{j}))_{j=1}^{m}$ using the data loss:

$$ $\mathcal{L}_{\text{DATA}}=\mathbb{E}_{\gamma,f,g}\left[||u_{\Theta_{L}}-u_{ \gamma,f,g}||\right].$ (10) $$

The expectation is computed on the distribution of the PDE parameters $(\gamma,f,g)$. The solution $u$ is entirely determined by these parameters as indicated by the notation $u_{\gamma,f,g}$. $||u_{\Theta_{L}}-u||$ denotes a distance between the target $(u(x_{j}))_{j=1}^{m}$ and the forecast $(u_{\Theta_{L}}(x_{j}))_{j=1}^{m}$ with $m$ the trajectory size. In practice, one samples a set of PDE instances $(\gamma_{i},f_{i},g_{i})$ and for each instance a corresponding sample $u_{i}$.

Theoretical analysis and relation to preconditioning
Analyzing the behavior of the inference algorithm is challenging due to the non linear nature of the solver. We however could get some intuition using simplifying assumptions. We build on the ideas introduced in section 2. We only sketch the main insights here and refer to appendix B for a proof and a more detailed analysis.

- •
Using a linearization of the neural solver it can be shown that the solver performs as a pre-conditioner on the linear system.
- •
Assuming that solution $u_{L}$ provided by the solver reaches the optimum $u^{*}$, and that the training set is such that the learned parameter $\Theta$ vectors span the whole parameter space of the model, then the convergence of the solver is guaranteed at an optimal rate.
- •
In practice, and as shown in the experiments, this suggests that the convergence rate is improved w.r.t. the reference baseline gradient algorithm. In the experiments section, we provide experimental comparisons with L-BFGS algorithm, which can be seen as a non linear pre-conditioning technique Rathore et al. ().

## 4 Experiments

We present the datasets used in the experiments in section 4.1, a comparison with selected baselines in section 4.2 and a test-time comparison with different optimizers demonstrating the remarkable effectiveness of the proposed method in section 4.3. More experiments can be found in the appendices: ablations are in appendix E and additional results and visualization in appendix F.

### 4.1 Datasets

We consider several representative parametric equations for our evaluation. More details about the data generation are presented in appendix D. Our objective is to learn a neural solver able to solve accurately and fast a new instance of a PDE, given its parametric form, and the values of the parameters $\gamma$, forcing terms $f$ and initial/boundary conditions $g$, i.e. $(\gamma,f,g)\mapsto u$. Solving is performed with a few iterations of the neural solver (algorithm 1).
For that, one trains the neural solver on a sample of the PDE parameter instances, see table 1 for the parameter distributions used for each parametric PDE.
Evaluation is performed on unseen sets of parameters within the same PDE family.

**Table 1: Parameters changed between each trajectory in the considered datasets.**
| Dataset | Parameters | Distribution |
| --- | --- | --- |
| Helmholtz | $\omega$ | $\mathcal{U}[0.5,50]$ |
| $u_{0}$ | $\mathcal{N}(0,1)$ |  |
| $v_{0}$ | $\mathcal{N}(0,1)$ |  |
| Poisson | $A_{i}$ | $\mathcal{U}[-100,100]$ |
| $u_{0}$ | $\mathcal{N}(0,1)$ |  |
| $v_{0}$ | $\mathcal{N}(0,1)$ |  |
| NLRD | $\nu$ | $\mathcal{U}[1,5]$ |
| $\rho$ | $\mathcal{U}[-5,5]$ |  |
| Darcy | $a(x)$ | $\psi_{\#}\mathcal{N}(0,(-\Delta+9I)^{-2})$ |
| with $\psi=12*\mathds{1}_{\mathbb{R}_{+}}+3*\mathds{1}_{\mathbb{R}_{+}}$ |  |  |
| Heat | $\nu$ | $\mathcal{U}[2\times 10^{-3},2\times 10^{-2}]$ |
| $J_{max}$ | {1, 2, 3, 4, 5} |  |
| $A$ | $\mathcal{U}[0.5,-0.5]$ |  |
| $K_{x}$, $K_{y}$ | $\{1,2,3\}$ |  |
| $\phi$ | $\mathcal{U}[0,2\pi]$ |  |

Helmholtz: We generate a dataset following the $1d$ static Helmholtz equation $u^{\prime\prime}(x)+\omega^{2}u(x)=0$ with boundary conditions $u(0)=u_{0}\text{ and }u^{\prime}(0)=v_{0}$. We generate $1,024$ trajectories with varying $\omega$, $u_{0}$ and $v_{0}$ with a spatial resolution of $256$.
Poisson: We generate a dataset following the $1d$ static Poisson equation with forcing term: $-u^{\prime\prime}(x)=f(x)$ with $u(0)=u_{0}$ and $u^{\prime}(0)=v_{0}$. The forcing term $f$ is a periodic function, $f(x)=\frac{\pi}{K}\sum_{i=1}^{K}a_{i}i^{2r}\sin(\pi x)$, with $K=16$ and $r=-0.5$. We generate $1,000$ trajectories with varying $u_{0},v_{0}$ and $f$ (through changing $a_{i}$) with a spatial resolution of $64$.
Reaction-Diffusion: In Krishnapriyan et al. (); Toloubidokhti et al. (), the authors propose a non-linear reaction-diffusion (NLRD). This PDE has been shown to be a failure case for PINNs (Krishnapriyan et al., ). We generate $1,000$ trajectories by varying the parameters of the PDE: $\nu$ and $\rho$ (see table 1). Spatial resolution is $256$ and temporal resolution is $100$. The PDE is solved on $[0,1]^{2}$.
Darcy Flow: The $2d$ Darcy Flow dataset is taken from (Li et al., ) and commonly used in the operator learning literature (Li et al., ; Goswami et al., ). For this dataset, the forcing term $f$ is kept constant $f=1$, and $a(x)$ is a piece-wise constant diffusion coefficient taken from (Li et al., ). We kept $1,000$ trajectories (on the $5,000$ available) with a spatial resolution is $64\times 64$.
Heat: The $2d+t$ Heat equation is simulated as proposed in (Zhou & Farimani, ). For this dataset, the parameter $\nu$ is sampled from $\mathcal{U}[2\times 10^{-3},2\times 10^{-2}]$ and initial conditions are a combination of sine functions with varying number of terms, amplitude and phase.
A summary of the datasets and the variable parameters for each dataset are presented in table 1 and more details on the dataset are provided in appendix D.
Experiments have been conducted on NVIDIA TITAN V ($12$ Go) for $1d$ datasets and NVIDIA RTX A6000 GPU with 49Go for $1d$ + time or $2d$ datasets. For all datasets, $800$ PDEs are considered during training and $200$ for testing. All metrics reported are evaluated on test samples (i.e. PDEs not seen during training. Coefficient as well as initial and/or boundary condition can vary from training).

### 4.2 Comparison with baselines

We performed comparisons with several baselines including fully data-driven supervised approaches trained from a data-loss only, unsupervised methods relying only on a PDE loss, and hybrid techniques trained from PDE + DATA loss. Networks size and training details are described in appendix C. In this experiment, we considered training the models using the training sets (physical losses or MSE when possible) unless stated otherwise.

Fully-supervised
We train a standard MLP to learn the mapping $(\gamma,f,g)\mapsto\Theta$, using as loss function $\mathcal{L}_{\text{DATA}}=\mathbb{E}_{\gamma,f,g,u}\left[||u_{\Theta_{L}}-u_{
\gamma,f,g,u}||\right]$
with $u_{\Theta}(x)=\sum_{i=0}^{N}\theta_{i}\psi_{i}(x)$, the $\psi_{i}(.)$ being fixed B-Spline basis functions (see appendix C). We denote this baseline as MLP+basis. Unsupervised
We compare our approach with unsupervised physics-informed models (Raissi et al., ). While the initial version of PINNs solves only one PDE instance at a time and requires retraining for each new instance, we developed here a parametric version of PINNs (PPINNs) where the PDE parameters are fed to the network. A similar approach has been used by (Zhang et al., ). In addition to PINNs, we also compare our method to the Physic-informed DeepONet (designed as PO-DeepONet for Physics-Only DeepONet) (Wang et al., ), which is designed to learn an operator for function-to-function mappings from physical losses and handles parametric PDEs. The mapping learned for the two unsupervised baselines is $(x,\gamma,f,g)\mapsto u(x)$. In order to provide a fair comparison with our optimization method, we fine-tuned the unsupervised baselines for each specific PDE instance for a few steps ($10$ or $20$). Comparison to pre-conditionning We compare our approach with vanilla PINNs (Raissi et al., ), i.e. by fitting one PINN per PDE in the test set and averaging. We optimize the PDE losses using L-BFGS (Liu & Nocedal, ) and refer to this baselines as PINNs+L-BFGS. As discussed in (Rathore et al., ), L-BFGS can be considered as a non linear preconditioning method for the linear Physics-Informed methods and fastens convergence. For this baseline one model is trained and evaluated for each PDE in the test set. We report the reader to appendix C for more details on the training procedure.
Hybrid
Finally, we compare with neural operators, i.e., trained to learn mappings $(x,\gamma,f,g)\mapsto u(x)$ using a combination of physical and data loss: $\mathcal{L}_{\textnormal{DATA}}+\mathcal{L}_{\textnormal{PDE}}$. We use as baselines Physics-Informed Neural Operator (PINO) (Li et al., ) and Physics-Informed DeepONet (PI-DeepONet) (Goswami et al., ). As already indicated, for a fair comparison, the Unsupervised and Hybrid baselines, are fine-tuned on each specific PDE instance for a few steps ($10$ on all dataset except for Heat for which $20$ steps are made).
Ours
We represent the solution $u_{\Theta}$ with a linear combination of B-Spline functions for $\Psi$ (Piegl & Tiller, ). This was motivated by the nature of B-Splines which allows to capture local phenomena. However, other bases could be used such as Fourier, Wavelet or Chebychev Polynomials. More details about the construction of the B-Spline basis are available in appendix C. The neural solver $\mathcal{F}_{\varrho}$ is composed of Fourier Layers (Li et al., ) that allow us capturing the range of frequencies present in the phenomenon.

The proposed method and all the baselines are implemented with PyTorch (Ansel et al., ).

**Table 2: Results of trained models - metrics in Relative MSE on test set. Best performances are highlighted in bold, second best are underlined.**
|  |  | 1d | 1d+time | 2d | 2d+time |  |
| --- | --- | --- | --- | --- | --- | --- |
|  | Baseline | Helmholtz | Poisson | NLRD | Darcy-Flow | Heat |
| Supervised | MLP + basis | 4.66e-2 | 1.50e-1 | 2.85e-4 | 3.56e-2 | 6.00e-1 |
| Unsupervised | PPINNs | 9.89e-1 | 4.30e-2 | 3.94e-1 | 8.47e-1 | 1.27e-1 |
| PINNs+L-BFGS | 9.86e-1 | 8.83e-1 | 6.13e-1 | 9.99e-1 | 9.56e-1 |  |
| PO-DeepONet | 9.83e-1 | 1.43e-1 | 4.10e-1 | 8.33e-1 | 4.43e-1 |  |
| Hybrid | PI-DeepONet | 9.79e-1 | 1.20e-1 | 7.90e-2 | 2.76e-1 | 9.18e-1 |
| PINO | 9.99e-1 | 2.80e-3 | 4.21e-4 | 1.01e-1 | 9.09e-3 |  |
| Neural Solver | Ours | 2.41e-2 | 5.56e-5 | 2.91e-4 | 1.87e-2 | 2.31e-3 |

##### Results:

Table 2 presents the comparison with the baselines. Each method is trained to solve a parametric equation, using a sample of the equation parameters. The evaluation set is composed of several PDE instances sampled from unseen PDE parameters $(\gamma,f,g)$.
The proposed method is ranked first or second on all the evaluations. The most comparable baselines are the unsupervised methods, since at inference they leverage only the PDE residual loss, as our method does. Therefore our method should be primarily compared toe these baselines. Supervised and hybrid methods both incorporate a data loss and make different assumptions while solving a different optimization problem.

Table 2 clearly illustrates that unsupervised Physics-informed baselines all suffer from ill-conditioning and do not capture the dynamics. Compared to these baselines, the proposed method improves at least by one order of magnitude on all cases. This is observed for PPINNs as well as on PINNs fitted on one equation only.
PINNs baseline performs poorly on these datasets because of the ill-conditioning nature of the PDE, requiring numerous optimization steps to achieve accurate solutions (appendix C). We observe that our neural solver has better convergence properties than other Physics-Informed methods. At will be seen later it also converges much faster.

The supervised baseline performs well on all the PDEs except Poisson and Heat. The data loss used for training this model is the mean square error which is well-behaved and does not suffer from optimization problems as the PDE loss does. We note that our method reaches similar or better performances on every dataset, while relying only on physical information at inference (algorithm 1) and solving a more complex optimization problem.

The hybrid approaches, do not perform well despite taking benefits from the PDE+DATA loss and from adaptation steps at test time. Again, the proposed method is often one order of magnitude better than the hybrids except on NLRD, where it has comparable performances. This shows that the combination of physics and data losses is also hard to optimize suffering from ill-conditioning.

### 4.3 Optimization for solving new equations

The main motivation for our learned PDE solver is to accelerate the convergence to a solution, w.r.t. predefined solvers, for a new equation. In order to assess this property, we compare the convergence speed at test time inference with classical solvers, PINNs and pre-trained PINO as detailed below. Results are presented in fig. 2 for the Poisson equation with performance averaged on $20$ new instances of the Poisson equation. The experiments for the other datasets are provided in appendix F.

Figure: Figure 2: Test-time optimization based on the physical residual loss $\mathcal{L}_{\textnormal{PDE}}$ for new PDE on Poisson.
Refer to caption: x2.png

Baseline optimizers
As for the classical optimizer baselines, we used SGD, Adam (Kingma & Ba, ) and L-BFGS (Krishnapriyan et al., ). These optimizers are used to learn the coefficient of the B-Spline basis expansion in the model $u_{\Theta}(x)=\sum_{i=0}^{N}\theta_{i}\psi_{i}(x)$. This provides a direct comparison to our iterative neural solver.
PINNs - We also compare to the standard PINNs (Raissi et al., ), i.e. by fitting one Neural Network (NN) per equation. Note that this requires full training from scratch for each new equation instance and this is considerably more computationally demanding than solving directly the parametric setting. This baseline is similar to the Adam optimizer mentioned above, except that the ansatz for this experience is a multilayer perceptron instead of a linear combination of B-Spline basis.
Hybrid pre-training strategies -
Finally we compare against the hybrid PINO pre-trained on a set of several parametric PDE instances and then fine-tuned on a new PDE instance using only the PDE loss associated to this instance.
Ours - We train our model as explained in algorithm 2 and show here the optimization process at test-time. In order to perform its optimization, our model leverages the gradient of the physical loss and the PDE parameters (coefficients, initial/boundary conditions). In this experiment, we use $L=5$ steps for a better visualization (whereas, we used $L=2$ in table 2).

Result: fig. 2 compares the number of optimization steps required for the different methods. Our neural solver converges very fast in only a few steps (5 here) to a good approximation of the solution, while all the other methods require thousands of iterations - we stopped here at $10,000$ steps. The classical optimizers (SGD, Adam, L-BFGS) do not converge for a new equation. The baseline PINNs trained here from scratch on each new equation shows an erratic convergence behavior. Pretrained PINO behaves better than the other baselines but still did not converged after $10,000$ steps. This clearly demonstrates the potential of our learned solver to deal with physical losses w.r.t. alternative pre-defined solvers.

## 5 Conclusion

We have presented a learning based PDE solver that allows fast test-time optimization of physical losses. Our method succeeds to considerably accelerate the optimization process for the complex problem of minimizing Physical losses, and is several orders of magnitude faster than classical hand-defined optimization methods such as Adam or L-BFGS.

##### Limitations and Future Work

The proposed method although efficient, could still be improved in several ways leading to further exploration of this idea. First, the memory needed to train our iterative algorithms can be more demanding than for usual ML model learning. Indeed, back-propagating through the iterations adds complexity and this might become an issue on higher dimensional basis. Second, we focused here on solution approximations expressed as a linear expansion in a pre-defined bases. More expressive representations such as neural networks could be investigated, although our preliminary experiments have shown even more ill-conditioning, due to the compositional nature of the neural network. More sophisticated training schemes, optimizers, or initializations could enhance the performances of the optimization process.
Future work will explore such issues to improve its scalability and capabilities for broader application of the proposed method.

##### Ethics Statement

Solving PDE is of crucial interest in many applications of science and engineering. While we do not directly target such real-world applications in this paper, one should acknowledge that solvers can be used in various ranges of scenarios including weather, climate, medical, aerodynamics, and military applications.

##### Reproducibility Statement

We will provide all elements that can be necessary to reproduce the experiments. Hyper-parameters, baselines configurations and training details are shown in appendices C, 4, 3, 2 and 1. The generation of the datasets are explained in appendices D, 1 and 5 with the solver used or the source, as well as the parameter ranges considered in the experiments. Finally, we provide a theoretical analysis of the model, under ideal scenario in appendix B. Code will be released upon acceptance as well as scripts to generate data when needed.

#### Acknowledgments

We acknowledge the financial support provided by DL4CLIM (ANR-19-CHIA-0018-01), DEEPNUM (ANR-21-CE23-0017-02), PHLUSIM (ANR-23-CE23-0025-02), and PEPR Sharp (ANR-23-PEIA-0008, ANR, FRANCE 2030).

## Appendix A Detailed related Work

##### PDE solvers:

Many tools for numerically solving PDEs have been developed for years. The standard methods for PDE include Finite Differences (FD), Finite Volume (FV), Finite Element Method (FEM), spectral and multigrid methods, and many others (S. H, ; Liu, ). While these methods are widely used, they often suffer from a high computational cost for complex problems or high-precision simulations. To address these challenges, integrating deep learning (DL) into solvers has emerged as a promising approach. Current solutions include incorporating correction terms into mathematical solvers to reduce numerical errors (Um et al., ). Some work such as Hsieh et al. (); Li et al. (); Rudikov et al. (); Kopaničáková et al. () build a method to directly enhance the convergence of numerical solvers through preconditioner learning. As an example, Rudikov et al. (); Li et al. () uses a neural operator to approximate conditioner for the flexible conjugate gradient method or Hsieh et al. () for the Jacobi method. Another example of preconditioner learning lies in Li et al. (), where the author uses GNN to assess symmetry and positive definiteness.

##### Unsupervised training:

Physics-Informed Neural Networks (PINNs) (Raissi et al., ) have been a pioneering work in the development of DL method for physics. In these models, the solution is a neural network that is optimized using the residual loss of the PDE being solved.
However, this method suffers from several drawbacks. First, as formulated in Raissi et al. (), PINNs can solve one instance of an equation at a time. Any small change in the parameters of the PDE involves a full retraining of the network. Efforts such as Beltrán-Pulido et al. (); Zhang et al. () have attempted to address this limitation by introducing parametric versions of PINNs capable of handling parametric equations, while Huang et al. (); Qin et al. () explores meta-learning approaches. Other approaches to improve PINNs generalization include using neural operator (Wang et al., ) or hyper-network (de Avila Belbute-Peres et al., ).
Moreover, PINNs have shown convergence difficulties: Krishnapriyan et al. () show that PINNs’ losses have complex optimization landscapes, complicating training despite adequate neural network expressiveness. Approaches like those detailed in Wang et al. () adopt a Neural Tangent Kernel (NTK) perspective to identify reasons for failure and suggest using adaptive weights during training to enhance performance. Additionally, studies such as Ryck et al. () demonstrate that PINNs suffer from ill-conditioned losses, resulting in slow convergence of gradient descent algorithms. Recently, (Rathore et al., ) show how training strategies can improve convergence of Physics-Informed Neural Networks and show that specific optimizers such as L-BFGS acts as conditioner on the physical losses.

##### Supervised training:

In contrast to the unsupervised training of Physics-informed Neural Networks, purely data-driven models have demonstrated remarkable capabilities for PDE simulation and forecasting. In most of the existing literature, the entire solver is replaced by a DL architecture and focuses on directly computing the solution from a given input data. A widely studied setting is operator learning which learns mappings between function spaces (Li et al., ; Kovachki et al., ; Lu et al., ). This method is very efficient, with the downside of relying on quite large quantities of data for training in order to ensure adequate generalization. Additionally, the neural network does not have access to the PDE in itself, only indirectly through the data.
To ensure physical constraint in purely-data-driven training, hybrid models have been proposed. The latter rely on both the available physical knowledge and some available data. Some examples include the Aphinity model (Yin et al., ) (where the authors have partial knowledge of the physics and learned the remaining dynamics from data), Physics-informed Deep Operator Networks (PIDON) (Wang et al., ; Goswami et al., ), Physics-informed Neural Operator (PINO) (Li et al., ) (DeepONet architecture (Lu et al., ) or Neural Operator models (Kovachki et al., ; Li et al., ) respectively with a combination of data and physical losses).

##### Learning to solve

:
Improving the learning scheme and optimizers through data-driven training has been studied since Li & Malik () and Andrychowicz et al. (). These works propose to learn the optimizer of neural networks, which are classically optimized through gradient-based algorithms such as Adam. They focus on improving training strategies for neural networks, which do not suffer from the optimization issues and ill-conditioning properties of physics-informed losses. We refer the reader to the survey of Chen et al. () for a complete overview. The closer work to ours is the very recent work of Bihlo () in which the author asseses the capabilities of learned optimizers for physics-informed neural networks. The main difference with our work relies on the problem setting. Bihlo () considers learning an optimizer on a single equation, and for different neural networks initialization, while we focus on efficiently solving several instances of parametric PDE with varying PDE parameters $\gamma,f,g$.

## Appendix B Proof of convergence

In this section, we propose a clarification on the convergence properties of Physics-Informed Neural Networks (PINNs) when a neural network is introduced into the framework. In section 2, we showed that gradient descent on the PINNs loss for a linear equation guarantees convergence, with the convergence rate governed by the condition number $\kappa(A)$. Here, we will address how incorporating a neural network impacts convergence, particularly focusing on the condition number of the system. We use the notations and assumptions introduced in section 2.

### B.1 Convergence in the Linear Case

In the linear case, the convergence of gradient descent is characterized by the number of steps $N(\epsilon)$ required to achieve an error less than $\epsilon$. This is linearly dependent on $\kappa(A)$ (section 2):

$$ $N(\epsilon)=O\left(\kappa(A)\log\left(\frac{1}{\epsilon}\right)\right)$ $$

Where:

- $A$: The matrix arising from the problem formulation in the basis.
- $\kappa(A)$: The condition number of $A$, which can be quite large, indicating a potential need for many iterations.

### B.2 Neural Network Approximation

When introducing a neural network $\mathcal{F}_{\varrho}$, we recognize the challenges in proving convergence due to the complex nature of neural networks. However, a useful approach is to consider the neural network’s behavior through its linear approximation, which allows us to use similar analysis techniques.

##### Linearization Approach

The neural network $\mathcal{F}_{\varrho}$ can be approximated by its tangent linear model, $\mathcal{F}_{\varrho}(v)\approx Pv$, where $P=\text{Jac}_{0}(\mathcal{F}_{\varrho})$, its Jacobian. The gradient descent update rule in this case becomes:

$$ $\Theta_{l+1}=\Theta_{l}-\eta P\nabla_{\Theta}\mathcal{L}_{\text{PDE}}(\Theta_{ l})=(\text{Id}-\eta PA)\Theta_{l}+c$ $$

Here, $c$ is a constant vector independent of $\Theta$.

##### Convergence Analysis with Linearization

By considering the linearized model, the updates can be expressed as:

$$ $\Theta_{l+1}=(\text{Id}-\eta PA)^{l+1}(\Theta_{0}-\Theta^{*})+\Theta^{*}$ $$

where:

- $\Theta^{*}=\Theta_{0}+A^{-1}c$ is the optimal parameter vector.

This expression parallels the convergence analysis of the original linear system, with the key difference being the introduction of the matrix product $PA$.

##### Condition Number of P ⁢ A 𝑃 𝐴 PA italic_P italic_A

In this formulation, convergence is governed by the condition number $\kappa(PA)$, rather than $\kappa(A)$:

$$ $N^{\prime}(\epsilon)=O\left(\kappa(PA)\log\left(\frac{1}{\epsilon}\right)\right)$ $$

The introduction of $P$ as a preconditioner often results in $\kappa(PA)<\kappa(A)$.
This implies that:

- •
the neural network can effectively reduces the condition number, enhancing convergence speed and stability.

We show below that the condition number is improved under some optimality conditions.

### B.3 Optimality and Convergence

The neural network $\mathcal{F}_{\varrho}$ is trained to minimize the $L_{2}$ norm $\|u_{L}-u^{*}\|_{2}$, where $u_{L}$ is the output after $L$ iterations and $u^{*}$ is the true solution. At convergence, $u_{L}=u^{*}$, implying $\Theta_{L}=\Theta^{*}$. We can consider the concatenated matrices $\Xi_{0}-\Xi^{*}$ where each row is associated to a sample. Incorporating this optimality condition, assuming the training samples span the parameter space completely (i.e. the matrix $(\Xi-\Xi^{*})^{\top}(\Xi-\Xi^{*})$ is full rank), we have that:

$$ $(\text{Id}-\eta PA)^{l+1}(\Xi_{0}-\Xi^{*})=0$ $$

With $\Xi_{0}\neq\Xi^{*}$, this condition implies that $(\text{Id}-\eta PA)^{l+1}=0$, indicating that $(\text{Id}-\eta PA)$ is nilpotent. As a consequence, all eigenvalues of $\eta PA$ are 1, leading to $\kappa(PA)=1$.

Under these conditions:

- •
This optimal condition number implies that convergence is not only guaranteed but also optimal, requiring fewer iterations.

∎

### B.4 Conclusion

As a conclusion of this proof sketch, we showed that, under optimality conditions and with linearization of the neural network, we have:

- (i)
Guaranteed Convergence: The method reliably achieves convergence to the optimal solution due to the reduced condition number.
- (ii)
Optimal Convergence Speed: With $\kappa(PA)=1$, the neural network provides an enhanced convergence rate, resulting in fewer required iterations compared to the original system without the neural network.

This analysis underscores the effectiveness of using neural networks within the Physics-Informed framework, leveraging their ability to reduce the condition number and improve convergence characteristics.

## Appendix C Implementation details

We add here more details about the implementation and experiments presented in section 4.

We implemented all experiments with PyTorch (Ansel et al., ). We estimate the computation time needed for development and the different experiments to be approximately $300$ days.

### C.1 B-Spline basis

We chose to use a B-Spline basis to construct the solution. We manually build the spline and compute its derivatives thanks to the formulation and algorithms proposed in (Piegl & Tiller, ). We used Splines of degree $d=3$ and constructed the Splines with $2$ different configurations:

- •
Take $N+d+1$ equispaced nodes of multiplicity $1$ from $\frac{d}{N}$ to $1+\frac{d}{N}$. This gives a smooth local basis with no discontinuities (see fig. 3(a)) represented by a shifted spline along the x-axis (denoted as shifted in the following).
- •
Use $N+1-d$ nodes of multiplicity $1$ and $2$ nodes of multiplicity $d$ (typically on the boundary nodes: 0 0 and $1$). This means that nodes 0 0 and $1$ are not differentiable (see fig. 3(b)). We call this set-up equispaced.

Figure: (a)
Refer to caption: extracted/5920248/img/bsplines_shifted.png

##### Higer-dimension basis

For $1$d + time, $2$d dataset and $2$d + time we build a $2$ (or $3$) dimensional B-spline basis, i.e., we treat the time coordinate as a variable as spatial ones. To build such bases, we compute the Cartesian product between the $2$ (or $3$) $1$d-bases, $1$ per dimension. This means that for a $2$d dataset, for which we chose to use bases of size $N1$ and $N2$ for each coordinate, the resulting basis will have $N1+N2+N1*N2$ terms. This method makes the training more costly and several techniques to improve its scalability could be used. For $3$d datasets, the number of terms in the basis is cubic.

### C.2 Training details

In our experiments, neural networks are trained using the Adam optimizer. For network optimization, we employ a smooth $l1$-loss for our solver while for other baselines, we use MSE loss and/or physical losses. All models are trained for at least $1,500$ epochs on datasets composed of some sampling of $\gamma$ and/or $f$ and/or $g$. If not stated otherwise, we train our model for $750$ epochs and baselines for $1,500$ epochs, with Adam and an initial learning rate of $0.001$. We use an exponential learning rate scheduler that lowers the learning rate by $0.995$ every epoch. Experiments are conducted on NVIDIA TITAN V ($12$ Go) for $1$d datasets and NVIDIA RTX A6000 GPU with 49Go for $1d$ + time, $2$d datasets or $2$d + time. We show on algorithm 3 a pseudo-code for training of our proposed method.

Figure: Algorithm 3 Training algorithm for learning to optimize physics-informed losses.

We make use of two nested components: the solver for providing the approximate solution to the PDE and the optimizer that conditions the training of the solver. Both are using gradient descent but with different inputs and objectives. The former optimizes the PDE loss (inner loop), while the latter optimizes the gradient steps of the solver through conditioning (outer loop). Optimization of the two components proceeds with an alternate optimization scheme. In particular, this implies that the map $F_{\varrho}$ is kept fixed during the inner optimization process.

### C.3 Models

We present here the architecture details for our model and the baseline on each dataset. In all our experiments, we use a GeLU (Hendrycks & Gimpel, ) activation function. The details of the model architecture on each dataset are presented in table 3.

**Table 3: Architecture details of our model and baselines**
| Model | Architecture | Dataset |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- |
| 1d | 1d + time | 2d | 2d + time |  |  |  |
| Helmholtz | Poisson | NLRD | Darcy | Heat |  |  |
| PINNs | MLP depth | 3 | 3 | 3 | 3 | 3 |
| MLP width | 256 | 256 | 256 | 256 | 256 |  |
| learning rate | 1e-4 | 1e-5 | 1e-4 | 1e-3 | 1e-3 |  |
| training steps | 1000 | 1000 | 1500 | 1500 | 1500 |  |
| PPINNs | MLP depth | 8 | 3 | 5 | 3 | 3 |
| MLP width | 64 | 256 | 256 | 256 | 256 |  |
| PO-DeepONet | Branch Net depth | 3 | 2 | 5 | 5 | 5 |
| Branch Net width | 256 | 256 | 256 | 256 | 256 |  |
| Trunk Net depth | 3 | 2 | 5 | 5 | 5 |  |
| Trunk Net width | 256 | 256 | 256 | 256 | 256 |  |
| MLP + basis | MLP depth | 5 | 5 | 5 | 5 | 5 |
| MLP width | 256 | 256 | 1,024 | 1,024 | 1,024 |  |
| PI-DeepONet | Branch Net depth | 3 | 2 | 5 | 5 | 5 |
| Branch Net width | 256 | 256 | 256 | 256 | 5 |  |
| Trunk Net depth | 3 | 2 | 5 | 5 | 5 |  |
| Trunk Net width | 256 | 256 | 256 | 256 | 256 |  |
| PINO | FNO depth | 3 | 3 | 4 | 3 | 3 |
| FNO width | 64 | 64 | 64 | 64 | 64 |  |
| FNO modes 1 | 16 | 16 | 10 | 20 | 7 |  |
| FNO modes 2 | - | - | 5 | 20 | 7 |  |
| FNO modes 3 | - | - | - | - | 5 |  |
| FNO fc dim | 64 | 64 | 64 | 64 | 64 |  |
| Ours | FNO depth | 3 | 3 | 3 | 3 | 3 |
| FNO width | 64 | 64 | 64 | 64 | 64 |  |
| FNO modes 1 | 16 | 16 | 10 | 20 | 7 |  |
| FNO modes 2 | - | - | 5 | 20 | 7 |  |
| FNO modes 3 | - | - | - | - | 5 |  |
| FNO fc dim | 64 | 64 | 64 | 64 | 64 |  |

##### Hyperparameters for Helmholtz:

For the baselines, we empirically searched hyperparameters to allow the network to handle the high frequencies involved in the Helmholtz dataset. Unfortunately, other network architectures did not improve the results. We proceed similarly for other datasets.

##### Training of PPINNs:

For the parametric PINN baseline, we considered different training schemes as using a unique optimizer or an improved training scheme from (Krishnapriyan et al., ), i.e. using mixed optimization with Adam and L-BFGS to train the model. For the Helmholtz dataset, we trained our model for 2000 epochs with a plateau learning rate scheduler with patience of 400 epochs. For the Poisson dataset, we trained our parametric PINN model for 5000 epochs with a cosine annealing scheduler with a maximum number of iterations of 1000. For the Reaction-diffusion dataset, we consider an initial learning rate of $0.0001$ instead of $0.001$.

##### PINNs baseline for Preconditioning:

The conditioning of the problem highly depends on the parameters of the PDE and initial/boundary conditions. This can lead to unstable training when optimizing PINNs, requiring specific parameters configuration for each PDEs. To avoid an extensive research of the best training strategy, we found a configuration that allowed a good fitting of most of the testing dataset. This means that in the values reported in table 2, we removed trainings for which the losses exploded (only a few hard PDEs were removed, typically between none to $20$). This prevents us from extensive hyper-parameter tuning on each PDE. Please note that this lower the reported relative MSE, thus advantaging the baseline. Finally, for computational reasons, we did not fitted the PINNs for more than $1,500$ epochs depending on the dataset. This can explain the observed performances of this baseline since PINNs are traditionally trained for more than $10,000$ epochs. We detail in table 4 the hyper-parameters for each dataset.

**Table 4: Hyper parameters for PINNs+L-BFGS baseline.**
| Hyper-parameter | Dataset |  |  |  |  |
| --- | --- | --- | --- | --- | --- |
| 1d | 1d + time | 2d | 2d + time |  |  |
| Helmholtz | Poisson | NLRD | Darcy | Heat |  |
| epochs | $1,000$ | $1,000$ | $1,500$ | $1,500$ | $1,000$ |
| learning rate | 1e-4 | 1e-5 | 1e-4 | 1e-3 | 1e-3 |

Moreover, we used a $3$ layer NN with $256$ neurons and a GeLU activation function.

##### Basis configuration:

For all datasets, we use Spline of degree $3$, built with shifted nodes. We change the number of terms in the basis depending on the problem considered. For the $1d$-problem, we use $32$ terms. For the $2d$-problem, $40$ elements are in the basis of each dimension, except for Reaction-Diffusion where the variable t has $20$ terms.
Moreover, during training, we use the projection of the initial conditions and/or parameters and/or forcings terms function in the basis as input to the networks.

## Appendix D Dataset details

For all datasets, we kept $800$ samples for training and $200$ as testing examples (except otherwise stated in the experiments).

### D.1 Helmholtz

We generate a dataset following the $1d$ static Helmholtz equation eq. 11.
For $x\in[0,1[$,

$$ $\begin{cases}u(x)^{\prime\prime}+\omega^{2}u(x)&=0,\\ u(0)&=u_{0},\\ u^{\prime}(0)&=v_{0}.\end{cases}$ (11) $$

The solution can be analytically derived: $u(x)=\alpha\cos(\omega x+\beta)$, with $\beta=\arctan(\frac{-v_{0}}{\omega u_{0}})$, $\alpha=\frac{u_{0}}{\cos(\beta)}$ and directly computed from the PDE data. We generate $1,024$ trajectories for training and $256$ for testing with $u_{0},v_{0}\sim\mathcal{N}(0,1)$ and $\omega\sim\mathcal{U}(0.5,50)$ and compute the solution on $[0,1]$ with a spatial resolution of $256$. For training, we keep $800$ samples and use the complete dataset for the additional experiments presented in figs. 12(a) and 12(b). Moreover, we sub-sample the spatial resolution by $4$ and keep $64$ points.

Figure: (a)
Refer to caption: extracted/5920248/img/vis_datasets_helmholtz-hf_0.png

### D.2 Poisson

We generate a dataset following the $1d$ static Poisson equation eq. 12 with forcing term. For $x\in[0,1[$,

$$ $\begin{cases}-u^{\prime\prime}(x)&=f(x),\\ u(0)&=u_{0},\\ u^{\prime}(0)&=v_{0}.\end{cases}$ (12) $$

We chose $f$ to be a non-linear forcing terms: $f(x)=\frac{\pi}{K}\sum_{i=1}^{K}a_{i}i^{2r}\sin(\pi x)$, with $a_{i}\sim\mathcal{U}(-100,100)$, we used $K=16$, $r=-0.5$, and solve the equation using a backward finite difference scheme. We generate $1,000$ trajectories with $u_{0},v_{0}\sim\mathcal{N}(0,1)$ and compute the solution on $[0,1]$ with a spatial resolution of $64$.

Figure: (a)
Refer to caption: extracted/5920248/img/vis_datasets_forcingmspoisson_1.png

##### Reaction-Diffusion

We use a non-linear reaction-diffusion used in (Krishnapriyan et al., ; Toloubidokhti et al., ). This PDE has been shown to be a failure case for PINNs (Krishnapriyan et al., ). The PDE states as follows:

$$ $\displaystyle\frac{\partial u(t,x)}{\partial t}-\nu\frac{\partial^{2}u(t,x)}{ \partial x^{2}}-\rho u(t,x)(1-u(t,x)$ $\displaystyle=0,$ (13) $\displaystyle u(0,x)=e^{-32(x-1/2)^{2}}.$ (14) $$

We generate $800$ trajectories by varying $\nu$ in $[1,5]$ and $\rho$ in $[-5,5]$. Spatial resolution is $256$ and temporal resolution is $100$, which we sub-sample by $4$ for training, leading to a spatial resolution of $64\times 25$. The PDE is solved on $[0,1]^{2}$ as in (Toloubidokhti et al., ).

Figure: (a)
Refer to caption: extracted/5920248/img/vis_datasets_1dnlrd_0.png

##### Reaction-Diffusion with initial conditions

: To complexify the setting, we also change initial condition of the problem (NLRDIC in the following). The initial condition express as follows:

$$ $u(x,0)=\sum_{i=1}^{3}a_{i}e^{-\frac{\left(\frac{x-h/4}{h}\right)^{2}}{4}}.$ (15) $$

Where $a_{i}$ are randomly chosen in $[0,1]$ and $h=1$ is the spatial resolution.

Figure: (a)
Refer to caption: extracted/5920248/img/vis_datasets_1dnlrdics_4.png

### D.3 Darcy Flow

The $2d$ Darcy Flow dataset is taken from (Li et al., ) and commonly used in the operator learning literature (Li et al., ; Goswami et al., ).

$$ $\displaystyle-\nabla.(a(x)\nabla u(x))$ $\displaystyle=f(x)\hskip 8.53581ptx\in(0,1)^{2},$ (16) $\displaystyle u(x)$ $\displaystyle=0\hskip 8.53581ptx\in\partial(0,1)^{2}.$ (17) $$

For this dataset, the forcing term $f$ is kept constant $f=1$, and $a(x)$ is a piece-wise constant diffusion coefficient taken from (Li et al., ). We kept $1,000$ trajectories (on the $5,000$ available) with a spatial resolution of $64\times 64$.

Figure: (a)
Refer to caption: extracted/5920248/img/vis_datasets_darcy_0.png

### D.4 Heat

As a proof that our method can handle $2$d + time, we consider the dataset proposed by (Zhou & Farimani, ).

$$ $\displaystyle\frac{\partial u(x,y,t)}{\partial t}-\nu\nabla^{2}u(x,y,t)$ $\displaystyle=0,$ (18) $\displaystyle u(x,y,0)$ $\displaystyle=\sum_{j=1}^{J}A_{j}\sin(\frac{2\pi l_{xj}x}{L}+\frac{2\pi l_{yj} y}{L}+\phi_{i}).$ (19) $$

Where $L=2$, $\nu$ is randomly chosen between $[2\times 10^{-3},2\times 10^{-2}]$, $A_{j}$ in $[-0.5,0.5]$, $l_{xj},l_{xy}$ are integers in $\{1,2,3\}$ and $\phi$ is in $[0,2\pi]$. As a difference with (Zhou & Farimani, ), we randomly chose $J$ between $1$ and $J_{max}=5$ to have more diversity in the represented frequencies in the data. The PDE are sampled with spatial resolution $64$ in x and y and temporal resolution $100$. However, during training these spatial resolution and subsampled by $4$ and coordinates are re scaled between 0 and 1. As for other PDEs, we use $800$ trajectories for training and $200$ for testing.

Figure: (a)
Refer to caption: x3.png

### D.5 Additional dataset: Advection

The dataset is taken from (Takamoto et al., ).

$$ $\displaystyle\frac{\partial u(t,x)}{\partial t}+\beta\frac{\partial u(t,x)}{ \partial x}$ $\displaystyle=0,\hskip 8.53581ptx\in(0,1),t\in(0,2],$ (20) $\displaystyle u(0,x)$ $\displaystyle=u_{0}(x),\hskip 8.53581ptx\in(0,1).$ (21) $$

Where $\beta$ is a constant advection speed, and the initial condition is $u_{0}(x)=\sum_{k_{i}=k_{1}...k_{N}}A_{i}\sin(k_{i}x+\phi_{i})$, with $k_{i}=\frac{2\pi{n_{i}}}{L_{x}}$ and ${n_{i}}$ are randomly selected in $[1,8]$. The author used $N=2$ for this PDE. Moreover, $A_{i}$ and $\phi_{i}$ are randomly selected in $[0,1]$ and $(0,2\pi)$ respectively. Finally, $L_{x}$ is the size of the domain (Takamoto et al., ).

The PDEBench’s Advection dataset is composed of several configurations of the parameter $\beta$ ($\{0.1,0.2,0.4,0.7,1,2,4,7\}$), each of them is composed of $10,000$ trajectories with varying initial conditions. From these datasets, we sampled a total of $1,000$ trajectories for $\beta\in\{0.2,0.4,0.7,1,2,4\}$ (which gives about $130$ trajectories for each $\beta$). This gives a dataset with different initial conditions and parameters. Moreover, during training, we sub-sampled the trajectories by $4$, leading to a grid of resolution $25$ for the t-coordinate and $256$ for the x-coordinate.

Figure: (a)
Refer to caption: extracted/5920248/img/vis_datasets_advections_0.png

### D.6 Summary of problem settings considered

A summary of the datasets and parameters changing between $2$ trajectories are presented in table 5.

**Table 5: Parameters changed between each trajectory in the considered datasets in the main part of the paper as well as additional datasets (Advections and NLRDIC).**
| Dataset | Changing PDE data | Range / Generation |
| --- | --- | --- |
| Helmholtz | $\omega$ | $[0.5,50]$ |
| $u_{0}$ | $\mathcal{N}(0,1)$ |  |
| $v_{0}$ | $\mathcal{N}(0,1)$ |  |
| Poisson | $A_{i}$ | $[-100,100]$ |
| $u_{0}$ | $\mathcal{N}(0,1)$ |  |
| $v_{0}$ | $\mathcal{N}(0,1)$ |  |
| Reaction-diffusion | $\nu$ | $[1,5]$ |
| $\rho$ | $[-5,5]$ |  |
| Darcy | $a(x)$ | $\psi_{\#}\mathcal{N}(0,(-\Delta+9I)^{-2})$ |
| with $\psi=12*\mathds{1}_{\mathbb{R}_{+}}+3*\mathds{1}_{\mathbb{R}_{+}}$ |  |  |
| Heat | $\nu$ | $[2\times 10^{-3},2\times 10^{-2}]$ |
| $J_{max}$ | {1, 2, 3, 4, 5} |  |
| $A$ | $[0.5,-0.5]$ |  |
| $K_{x}$, $K_{y}$ | $\{1,2,3\}$ |  |
| $\phi$ | $[0,2\pi]$ |  |
| Advection | $\beta$ | $\{0.2,0.4,0.7,1,2,4\}$ |
| $A_{i}$ | $[0,1]$ |  |
| $\phi_{i}$ | $[0,2\pi]$ |  |
| $k_{i}$ | $\{2k\pi\}_{k=1}^{8}$ |  |
| NLRDIC | $\nu$ | $[1,5]$ |
| $\rho$ | $[-5,5]$ |  |
| $a_{i}$ | $[0,1]$ |  |

## Appendix E Additional Experiments

### E.1 Ablation

We experimentally show some properties of our iterative method.
These evaluations are made on a test set composed of several instances of PDE with varying configurations $(\gamma,f,g)$. These experiments are performed on the Helmholtz equation which appeared as one of the more complex to optimize in our evaluation with $L=5$ unless stated otherwise.

##### Error w.r.t the number of steps L

On fig. 11, we show that having more optimizer steps allows for a better generalization. However, we observed in our experiments that the generalization error stabilizes or even increases after the proposed $5$ steps. This limitation of the solver should be investigated in future work in order to allow the model to make more iterations. Note that for this experiment, we lowered batch size (and adapted learning rate accordingly) when the number of steps increased (fig. 11).

Figure: Figure 11: Error on the test set (Helmholtz equation) w.r.t the number of iterations
Refer to caption: extracted/5920248/img/abL_tab.png

##### Error w.r.t the number of training samples

On figs. 12(a) and 12(b), we show that compared to other physics-informed baselines, our solver requires fewer data to learn to solve PDE. Note that on this simple example, the MLP+basis baseline also performed well. However, as shown in table 2, this is not the case for all the datasets.
The contribution of the iterative procedure clearly appears with $2\times$ less data, our models performs best than this baseline.

Figure: (a)
Refer to caption: x5.png

##### Optimization with ℒ PDE subscript ℒ PDE \mathcal{L}_{\textnormal{PDE}} caligraphic_L start_POSTSUBSCRIPT PDE end_POSTSUBSCRIPT

Table 6 shows that optimizing our network with physical loss greatly complexifies training. Indeed, adding an ill-conditioned loss to the standard MSE makes training ill-conditioned.

**Table 6: How the physical loss complexify training and lowers performances. Metrics in Relative MSE on the test set of Helmholtz equation.**
| Input | Relative MSE |
| --- | --- |
| $\mathcal{L}_{\textnormal{PDE}}$ + $\mathcal{L}_{\textnormal{DATA}}$ | 1.11 |
| $\mathcal{L}_{\textnormal{DATA}}$ | 2.19e-2 |

##### Iterative update & SGD-based update

In table 7, we compare two optimizer configurations. "Direct" means that the network directly predicts the parameters for the next step (i.e. $\theta^{k+1}=\mathcal{F}_{\varrho}(\nabla\mathcal{L}_{\textnormal{PDE}}(\theta
^{k}),\;\gamma,f,g)$), while "GD", corresponds to the update rule described before and using SGD as the base algorithm $\theta^{k+1}=\theta^{k}-\eta\mathcal{F}_{\varrho}(\nabla\mathcal{L}_{
\textnormal{PDE}}(\theta^{k}),\;\gamma,f,g)$. The latter clearly outperforms the direct approach and shows that leaning increments is more efficient than learning a direct mapping between two updates. As shown in table 7 increasing the number of steps improves the performance (shown here for 1 and 5 update steps). However the performance does not improve anymore after a few steps (not shown here).

**Table 7: Comparison of different optimizer configurations for solving the Helmholtz equation. Metrics in Relative MSE on the test set.**
|  | Relative MSE |  |
| --- | --- | --- |
| N-steps | Direct | GD |
| $1$-step | 1.08e-1 | 8.5e-2 |
| $5$-step | 9.07e-2 | 2.19e-2 |

##### Optimization with different inner learning rates

On table 8, we study the performance of our proposed method with different inner learning rates $\eta$. As expected, a higher learning rate leads to better performances since, the optimization is taking bigger steps.

**Table 8: Ablation on the inner learning rate. Metrics in Relative MSE on the test set of Helmholtz equation.**
| learning rate | Relative MSE |
| --- | --- |
| 0.01 | 7.32e-2 |
| 0.1 | 4.93e-2 |
| 1 | 2.19e-2 |

##### Quantifying the importance of input feature for the learned solver

As indicated in eq. 9 the inputs of our learned solver are $(\gamma,f,g,\nabla\mathcal{L}_{\textnormal{PDE}}(\theta_{L}))$.
We performed experiments by removing either $\gamma$ or $\nabla\mathcal{L}_{\textnormal{PDE}}(\theta_{l})$ from the input (For the Helmholtz equation, there is no forcing term $f$). The BC $g$ are kept since they are part of the PDE specification and are required for ensuring the uniqueness of the solution.
This experiment (table 9) illustrates that conditioning on the PDE parameters
$\gamma$ is indeed required to solve the parametric setting. Without $\gamma$, the solver has no hint on which instance should be used. The addition of the gradient information, $\nabla$, which is at the core of our method, is indeed crucial for improving the convergence and validates our setting.

**Table 9: Effect of using the gradient as input w.r.t the PDE parameters. Metrics in Relative MSE on the test set of Helmholtz equation.**
| Input | Relative MSE |
| --- | --- |
| $\gamma$ + g | 3.75e-1 |
| $\nabla$ + g | 1.07e-1 |
| $\nabla$ + $\gamma$ + g | 2.19e-2 |

Non-linear Reconstruction
Finally, we showed by a simple experiment that our model can handle non linear cases (see table 10). We propose to model the solution $u$ using a non-linear combination of the basis terms $\phi_{i}$. The relation between the $\phi_{i}$s is modeled using a simple NN with one hidden layer and a tanh activation function. This experiment is performed on the Poisson PDE.

**Table 10: Non linear combination of the basis. Relative MSE on the test set for our proposed method and comparison with other non linear models and optimizers.**
| baselines | Relative MSE |
| --- | --- |
| PINNs+L-BFGS | 8.83e-1 |
| PPINNs | 4.30e-2 |
| Ours | 3.44e-3 |

### E.2 Statistical analysis

We provide a statistical analysis on some datasets. These experiments are conducted using $L=5$ (instead of $L=2$ in table 2). For computational cost reasons, we did not make this experiment on all datasets, but kept $2$ datasets in $1d$ (table 11), 1 dataset in $1$d+time and 1 dataset in $2$d (table 12) so that several configurations and data sizes are represented.

**Table 11: Results of trained models with error bars on 1d datasets - metrics in Relative MSE on test set. Best performances are highlighted in bold, second best are underlined.**
|  |  | 1d |  |
| --- | --- | --- | --- |
|  | Baseline | Helmholtz | Poisson |
| Supervised | MLP + basis | 5.26e-2 $\pm$ 7.56e-3 | 1.58e-1 $\pm$ 7.98e-3 |
| Unsupervised | PPINNs | 8.33e-1$\pm$5.61e-3 | 3.59e-2$\pm$2.11e-2 |
| PO-DeepONet | 9.84e-1$\pm$6.93e-4 | 1.79e-1$\pm$3.09e-2 |  |
|  | PI-DeepONet | 9.81e-1 $\pm$ 2.25e-3 | 1.25e-1 $\pm$1.04e-2 |
| Hybrid | PINO | 9.95e-1 $\pm$ 3.30e-3 | 3.27e-3 $\pm$ 1.38e-3 |
|  | Ours | 2.17e-2 $\pm$ 1.12e-3 | 4.07e-5 $\pm$ 2.65e-5 |

**Table 12: Results of trained models with error bars on 1d + time and 2d datasets - metrics in Relative MSE on test set. Best performances are highlighted in bold, second best are underlined.**
|  |  | 1d + time | 2d |
| --- | --- | --- | --- |
|  | Baseline | 1dnlrd | Darcy-Flow |
| Supervised | MLP + basis | 2.83e-5 $\pm$ 6.83e-7 | 3.78e-2$\pm$2.09e-3 |
| Unsupervised | PPINNs | 4.64e-1 $\pm$ 1.92e-2 | 9.99e-1$\pm$2.63e-2 |
| PO-DeepONet | 4.18e-1 $\pm$ 1.04e-2 | 8.32e-1$\pm$2.51e-4 |  |
|  | PI-DeepONet | 7.88e-2 $\pm$ 1.96e-4 | 2.72e-1 $\pm$ 4.44e-3 |
| Hybrid | PINO | 8.00e-5 $\pm$ 1.00e-5 | 1.17e-1 $\pm$ 1.42e-2 |
|  | Ours | 2.61e-5 $\pm$ 2.53e-6 | 1.62e-2 $\pm$ 3.06e-4 |

### E.3 Computational cost

As we can see on table 13, our model takes longer to train due to the iterative process occurring at each epoch.
However, note that this is training time; inference time is similar to other methods (see table 14). We detail a justification for this additional training time compared to PINNs variants below:

- •
Comparison with vanilla PINNs:
Consider the following experiment. Suppose we train a classical PINN on a single instance of the Darcy PDE. Based on the training times shown in table 13, using this method, we performed $1,500$ steps, which took $420$ minutes for training (please note that the performances were less accurate than our model’s performance). If we wanted to train a PINNs on each PDE of out entire test dataset for $15000$ epochs (sometimes even more steps are required), this would take $4200$ minutes or stated otherwise approximatly $3$ days.
Suppose now, that one wants to solve an additional equation. This will require an average of $0.226$ seconds (see table 14) with our method, while PINNs would require an entirely new training session of approximately $20$ minutes for $15,000$ steps. This makes our method $5,000$ times faster than traditional PINNs for solving any new equation.
- •
Comparison with PINNs parametric variants: Now let us consider a parametric variant of PINNs to handle multiple PDEs (PPINNs for parametric PINNs). In table 2, PPINNs was trained for $5000$ epochs only. Let us consider training it further as suggested for vanilla PINNs, until $15,000$ epochs. First, this training would require approximately $19$h$30$m. Then, the optimization problem is still ill-conditioned, training further would probably not significantly improve the performance.

**Table 13: Training time of the experiments shown in Table 2. of the paper on a NVIDIA TITAN RTX (25 Go) GPU. d stands for days, h for hours, m for minutes.**
| Dataset | Helmholtz | Poisson | 1dnlrd | Darcy | Heat |
| --- | --- | --- | --- | --- | --- |
| MLP + basis | 30m | 20m | 1h10m | 2h | 4h45 |
| PPINNs | 15m | 20m | 4h15m | 6h30m | 1d2h |
| PINNs+L-BFGS | 15h15 | 7h35 | 20h30 | 7h | 13h |
| PODON | 10m | 10m | 3h30m | 1d9h | 22h |
| PIDON | 10m | 10m | 3h30m | 1d10h30m | 22h |
| PINO | 15m | 10m | 1h10m | 45m | 2h40 |
| Ours | 30m | 30m | 4h30 | 10h15 | 1d 13h |

**Table 14: Inference time on the test set (in seconds). All experiment are conducted on a single NVIDIA RTX A6000 (48Go). We report the mean time computed on the test set to evaluate the baselines as performed on Table 2 in the paper (i.e. with $10$ test-time optimization steps when applicable and $20$ steps on the Heat dataset).**
| Dataset | Helmholtz | Poisson | 1dnlrd | Darcy | Heat |
| --- | --- | --- | --- | --- | --- |
| MLP + basis | 1.12e-2 | 1.18e-2 | 1.25e-2 | 1.19e-2 | 1.66e-2 |
| PPINNs | 3.09e-1 | 2.03e-1 | 2.91e-1 | 3.22e-1 | 5.45e-1 |
| PODON | 3.27e-1 | 2.71e-1 | 4.38e-1 | 6.32e-1 | 8.85e-1 |
| PIDON | 3.32e-1 | 2.96e-1 | 4.43e-1 | 6.35e-1 | 8.80e-1 |
| PINO | 3.14e-1 | 1.24e-1 | 5.19e-1 | 2.21e-1 | 8.08e-1 |
| Ours | 2.58e-1 | 2.16e-1 | 2.84e-1 | 2.26e-1 | 2.90e-1 |

### E.4 Additional datasets

We provide additional experiment on $2$ new datasets: Non-linear Reaction-Diffusion in a more complex setting and Advections. We refer the reader to appendix D for the details about the PDE setting. These datasets were not included in the main part of the paper due to a lack of space in the results table.

**Table 15: Results of trained models on additional datasets - metrics in Relative MSE on test set. Best performances are highlighted in bold, second best are underlined.**
|  | Baseline | NLRDIC | Advections |
| --- | --- | --- | --- |
| Supervised | MLP + basis | 9.88e-4 | 6.90e-2 |
| Unsupervised | PPINNs | 3.71e-1 | 4.50e-1 |
| PO-DeepONet | 4.36e-1 | 5.65e-1 |  |
|  | PI-DeepONet | 5.39e-2 | 4.26e-1 |
| Hybrid | PINO | 3.79e-3 | 6.51e-4 |
|  | Ours | 1.41e-3 | 5.39e-3 |

These additional dataset show cases where baselines are performing very well on the considered PDE. For Advections, PINO reached very good performances. We believe that the Fourier layers used in the model fits well to the phenomenon. Indeed, the solution is represented using a combination of sine moving with time. This simple dynamics is easily captured using Fourier transformations. Our B-Splines basis can be sub-optimal for this dataset. For the complexified version of NLRD, NLRDIC, it is the supervised baseline that performs best. This dataset, does not present high frequencies, and the MLP looks sufficiently expressive to find the coefficient in the basis. However, this baseline had difficulties to reconstruct the higher frequencies in Advections.
Even if our model does not perform best on these datasets, it is ranked second. We believe that these additional results show the robustness of our method.

## Appendix F Qualitative results

This section is dedicated to visualization of the results of our model, baselines and optimizers, presented in section 4.
For each PDE considered, we chose $2$ samples in the test sets and compute the solutions with our model and the different baselines. We provide visualization samples with $L=5$ i.e. results proposed in tables 11 and 12 to detail more precisely the evolution of the solution at several steps of optimization. $3$ datasets are shown with $L=2$: Heat for computational reasons (training with $L=5$ is much more expensive when the dimension of the problem increases) and the $2$ additional datasets trained only with $L=2$ (Advections and NLRDIC). Then, we show the evolution of the reconstruction of the solution with our method i.e. we plot the solution at each step of the optimization (figs. 13(a), 13(b), 16(a), 16(b), 19(a), 19(b), 22(a), 22(b), 25(a), 25(b), 28(a), 28(b), 31(a) and 31(b)) and we compare the final prediction with baselines’ (figs. 14(a), 14(b), 17(a), 17(b), 20(a), 20(b), 23(a), 23(b), 26(a), 26(b), 29(a), 29(b), 32(a) and 32(b)). Finally, we chose $20$
($6$ for Heat) PDEs and we reproduce fig. 2 for every dataset. More precisely, we optimize one PINN per PDE using Adam, we fit our basis using several optimizers (GD, ADAM, L-BFGS and our learned optimization process) and we fine-tune the learned PINO for $10,000$ steps and visualize the evolution of the MSE (averaged at each step on the selected PDE). These figures show the relevance of learning the optimizer when using physics-informed losses (figs. 15, 18, 21, 24, 27, 30 and 33).

### F.1 Helmholtz

Figure: (a)
Refer to caption: x7.png

Figure: (a)
Refer to caption: x9.png

Figure: Figure 15: Test-time optimization based on the physical residual loss $\mathcal{L}_{\textnormal{PDE}}$ on Helmholtz. Note that, even though hardly visible on this figure, the optimization is running very slowly and the PINN MSE (orange) decreases for the last steps. This dataset will probably need even more steps before convergence.
Refer to caption: x11.png

### F.2 Poisson

Figure: (a)
Refer to caption: x12.png

Figure: (a)
Refer to caption: x14.png

Figure: Figure 18: Test-time optimization based on the physical residual loss $\mathcal{L}_{\textnormal{PDE}}$ on Poisson.
Refer to caption: x16.png

### F.3 Reaction-Diffusion

Figure: (a)
Refer to caption: x17.png

Figure: (a)
Refer to caption: x19.png

Figure: Figure 21: Test-time optimization based on the physical residual loss $\mathcal{L}_{\textnormal{PDE}}$ on NLRD.
Refer to caption: x21.png

### F.4 Darcy

Figure: (a)
Refer to caption: x22.png

Figure: (a)
Refer to caption: x24.png

Figure: Figure 24: Test-time optimization based on the physical residual loss $\mathcal{L}_{\textnormal{PDE}}$ on Darcy.
Refer to caption: x26.png

### F.5 Heat

Figure: (a)
Refer to caption: x27.png

Figure: (a)
Refer to caption: x29.png

Figure: Figure 27: Test-time optimization based on the physical residual loss $\mathcal{L}_{\textnormal{PDE}}$ on Heat. For computational reasons, this experiment has been conducted on $1,000$ steps only.
Refer to caption: x31.png

### F.6 Advection

Figure: (a)
Refer to caption: x32.png

Figure: (a)
Refer to caption: x34.png

Figure: Figure 30: Test-time optimization based on the physical residual loss $\mathcal{L}_{\textnormal{PDE}}$ on Advection.
Refer to caption: x36.png

### F.7 Non-Linear Reaction-Diffusion with Initial Conditions

Figure: (a)
Refer to caption: x37.png

Figure: (a)
Refer to caption: x39.png

Figure: Figure 33: Test-time optimization based on the physical residual loss $\mathcal{L}_{\textnormal{PDE}}$ on NLRD with varying IC.
Refer to caption: x41.png

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 3
presentation: 2
questions: -	Do you have any intuition on what the model $\mathcal{F}_{\rho}$ is learning? The paper interprets it as a preconditioner, but this answer is very vague and lack of information.

Think this way: given infinite data and a perfect training, what is the resulting mapping of $\mathcal{F}_{\rho}$?

Then consider: what is the effect of $\mathcal{F}_{\rho}$? Does it necessarily (or, ideally) help reduce the condition number?

-	Are the baseline hybrid methods trained with exactly the same dataset (and same amount of data) as the proposed method?

-	L308, isn’t $m$ the number of grid points instead of trajectories?
soundness: 2
strengths: -	The method is full of imaginations!

-	The empirical result in the test cases investigated are surprising. The proposed method help significantly accelerate the convergence of pde-loss optimization.
summary: This paper proposes a method to overcome the challenging optimization of PDE-loss in physics-informed methods by learning a modification on the gradient (w.r.t. model parameters), which intuitively serves as a preconditioner. The experiment displays a noticeable acceleration in convergence.
weaknesses: -	The theoretical analysis (main text and appendix B) is superficial and should not be considered as ‘proof’. In L814, the only important part of convergence analysis is circumvented by saying ‘the introduction of P as a preconditioner often results in κ(PA)<κ(A)’ instead of showing it. Thus the author actually didn’t prove anything. In fact, the effeteness of the method is highly related to whether it can help mitigate the condition number, which might not hold true, not supported by either theoretical or empirical evidence.
Also, L790-791 simplify the problem so much that the derivation could no longer provide insight to understanding the proposed method, unless the authors demonstrate the approximation to $Pv$ is reasonable in some sense.

-	To train the model $\mathcal{F}_{\rho}$, one has to first fix a set of grid point. This is neither flexible nor suitable to leverage the advantageous of PINN. In many PINN papers, the pde residual are estimated in a Monte-Carlo manner by drawing random points and computing its pointwise residual, so as to eliminate the heave computational and memory cost. Moreover, since the paper are target at solving a family of PDEs, it is always the case that one has to use finer grids / higher resolution to adequately capture small-scale effects for some harder instances in the large PDE class.

-	Regarding the method itself, in my opinion, the performance significantly depends on how well $\mathcal{F}_{\rho}$ 

generalizes to unseen or even out-of-distribution data. Intuitively it seems like it would be harder for $\mathcal{F}_{\rho}$ 

to generalize than the surrogate model (as in PINO or physics-DeepONet) itself, since the mapping $\mathcal{F}_{\rho}$ try 

to approximate is more complicate. A systematic study on how well $\mathcal{F}_{\rho}$ needs to be learnt so as to ensure reasonable improvement on pde-loss convergence is necessary to make this method practical.

-	As mentioned in the limitations, the method will possibly suffer from a memory issue and is slow to train since it needs to back-propagate through several optimization iterations. This  hinders the practicability of the method.

-	The PDE considered are either 1D or linear, which are all quite toyish. The paper could benefit from studying their methods in some equations will optimization indeed brings severe challenges for physics-informed methods, for instance Navier-Stokes with relatively high Reynolds numbers. Otherwise, the experiments are not convincing.

-	Many details important for evaluating this paper are only mentioned in the appendix, e.g. the training and inference time cost. I am aware of the page limit, but the content should be briefly summarized in the main text, otherwise few readers notice they are there.

To conclude, although the paper is rich in novelty and has impressive performance in the experiments shown in the paper, and has the potential to provide new insight to the community, it is still not fully convincing- more understanding, discussion, and evaluation are necessary. Therefore, though I appreciate the author’s effort, I vote for rejection.

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 2
questions: 1) Why did you not include wall-clock times?

2) Why did you not include a comparison(s) with a more sophisticated variant of PINN(s)?
soundness: 2
strengths: - This is the first time I've seen the use of learned optimizers in SciML. I think this approach has the potential to be very useful.
- The potential connection between learning the optimizer is interesting, though still somewhat tenuous; see weaknesses below.
- The proposed framework outperforms vanilla baselines, which is good.
summary: Physics-informed Neural Networks (PINNs) are popular neural network-based approaches for solving partial differential equations and provide a compelling alternative to traditional PDE solvers. Unfortunately, the PINN objective yields a difficult optimization problem to train the network. The difficulty stems from the presence of an ill-conditioned differential operator in the loss, which leads to an ill-conditioned objective. The present paper proposes a framework for solving parametric PDEs based on learning an optimizer tailored to that class of problems. The authors provide a heuristic argument showing that the learned optimizer can be viewed as preconditioning the original loss. Preliminary experiments show the proposed approach yields better performance relative to existing methods.
weaknesses: **Related work**

 While I have not seen the idea of learning an optimizer applied in the SciML literature, it certainly is not new in the overall ML literature. The authors should explain how their approach fits into this existing body of work in the related work section.
The authors should mention Cho et al.'s recent work, which proposed a very effective method for solving parametric PDEs using PINNs.

**Theory**

The theoretical support for the benefit of learning the optimizer is weak. 
The argument provided in Appendix B is heuristic and relies heavily on the assumption that both networks are well-approximated by their first-order Taylor expansions.
This regime only holds under certain strong assumptions, which the paper neglects to mention.
Finally, it is not apparent why it should be the case that $\kappa(PA) \ll \kappa(A)$. 
So, the potential connection between the learning optimizer is interesting but very tenuous.     
 
**Experiments**

The experiments in the paper are very weak. 
Comparing to the vanilla PINN alone using L-BFGS for training is inadequate. 
As is well known, better performance is achieved by using the combination of Adam+L-BFGS. 
Plus, more sophisticated optimizers, like Muller et al.'s natural gradient approach or Rathore et al.'s NNCG, can greatly enhance PINN training.
It would be much more compelling to compare to one of these more sophisticated approaches to training PINNs.

The paper mentions the ability to cover parametric PDEs as a contribution, but none of the experiments demonstrate this advantage. Moreover, the method introduced in this paper should be compared to a dedicated PINN method for parametric PDES like that of Cho et al. (2024).

The authors' approach clearly has a higher time cost than the vanilla approach, but they do not provide runtimes for their method relative to the vanilla method. This makes it challenging to develop a concrete idea of how much more expensive their approach is.

**Presentation**

The text has many grammatical errors, which can sometimes make the presentation difficult to follow. I recommend the authors do a careful read-through and make appropriate edits. 
This will greatly improve the readability of the paper.

**Overall** 

I believe the paper contains some interesting ideas that have the potential to enhance the applicability of PINNs. 
Unfortunately, the paper is not ready for publication at ICLR. 
The experiments are lacking, so it is difficult to determine if the approach here would improve over more sophisticated variants of PINNs,  which are used in practice.
In addition, the theoretical contribution is too heuristic, and the presentation needs polishing.

**References**

Cho, W., Jo, M., Lim, H., Lee, K., Lee, D., Hong, S. and Park, N., 2024. Parameterized physics-informed neural networks for parameterized PDEs. arXiv preprint arXiv:2408.09446.

Müller, J. and Zeinhofer, M., 2023, July. Achieving high accuracy with PINNs via energy natural gradient descent. In International Conference on Machine Learning (pp. 25471-25485). PMLR.

Rathore, P., Lei, W., Frangella, Z., Lu, L. and Udell, M., 2024. Challenges in training PINNs: A loss landscape perspective. arXiv preprint arXiv:2402.01868.

---

confidence: 3
contribution: 3
initial_rating: 8
presentation: 4
questions: My understanding is that the objective function for training the neural solver is data loss + PDE loss, which represents a soft-constrained optimization problem. Have you considered using a hard-constrained method, as in Lu et al. [1]?

[1] Lu, Lu, et al. "Physics-informed neural networks with hard constraints for inverse design." SIAM Journal on Scientific Computing 43.6 (2021): B1105-B1132.
soundness: 4
strengths: A very good paper demonstrating a novel approach to addressing PINN’s limitations, with extensive experimental results to support the claims. 

The paper presents strong motivation and a thorough review of related work. 

The methods and the overall paper are easy to follow.
summary: This paper proposes an optimization process that can learn a set of parametric PDEs quickly and precisely while addressing the ill-conditioning issues that traditional PINNs face. The authors first train a physics-informed neural network on a dataset containing different PDE coefficients, initial conditions, and boundary conditions, enabling the network to predict PDE solutions in the initial step. Then, a solution approximation expressed as a linear expansion in a pre-defined basis is optimized to adapt to each PDE instance. This method enables fast learning of different PDE instances and effectively overcomes ill-conditioning issues.
weaknesses: I have some questions about the training of the neural solver.

---

confidence: 3
contribution: 3
initial_rating: 8
presentation: 3
questions: My suggestions are included in the "Weaknesses" section and I would be very pleased to hear the authors' response to these suggestions.
soundness: 3
strengths: Overall the paper does a good job of explaining the problem and describing the proposed solution. However there are a number of issues that need to be expanded upon in order to improve the thoroughness and reliability of the findings. 

In general I think the paper represents a useful contribution to the literature and I am leaning towards an acceptance. However, I do think that it could be significantly improved if some of the findings detailed in the appendices could be brought into the main body of the paper as described in this review. Although the paper is fully 10 pages long, some of the descriptive text, especially in the early part of the paper could be made more concise to make space for these valuable aspects of the findings to be included in the main body of the paper.
summary: This paper addresses the problem of optimising Physics informed deep neural networks for solving PDE. Specifically, the paper addresses the issue of ill-conditioned optimization problems. 

Using a Fourier representation, the authors clearly lay out the problem with ill conditioning of general PINNs and give a numerical example to illustrate this. 

The paper describes how the problem can be transformed using a neural solver into one that can be optimised using gradient descent more quickly than the original problem.
weaknesses: The parameters $\Theta$ are updated using gradient descent. This is nested within a additional optimisation process for tuning the parameters $\varrho$. Some results should be included in the main body of the paper which show the CPU time, including for the optimisation of $\mathcal{F}_\varrho$. There are some results presented in Appendix F but these should be discussed in the main paper as they are important for evaluating the usefulness of the method presented. 

The paper describes the problem in terms of the condition number $\kappa(A)$. A sketch proof is included in Appendix B to show that convergence is guaranteed and its speed is enhanced using this structure. However the wording around this sketch proof is quite confusing. Section B.2 there states "The introduction of $P$ as a pre-conditioner \textit{often} results in $\kappa(PA)<\kappa(A)$." and "the condition number is improved under \textit{some} optimality conditions. However in Sections B.3 and B.4 convergence and optimality is described as \textit{optimal} and \textit{guaranteed}. This inconsistency should be addressed due to the importance of this in supporting the empirical findings within the main part of the paper. 

The results in the paper show very good improvements in the optimisation time and the MSE for the example PDEs used. However a range of parameters for each PDE is quoted in table 1. In table 2 the MSE results for the different equations are shown, however there is no indication of how these errors vary within the parameter ranges. There are some results presented  for this in Appendix E, however these should be brought into the main part of the paper as it is important to show the robustness of the method. 

In general the paper is well written and formatted but  In line 436 and 437, the sentence reads ``Therefore our method should be primarily compared toe these baselines."

---

confidence: 5
contribution: 1
initial_rating: 3
presentation: 2
questions: mentioned above.
soundness: 1
strengths: The authors tackle and important problem in AI and science.
summary: The paper propose a method that seems to be a PINO method to tackle learning of PDEs with the presents supervision from both data and equations.
weaknesses: The writing and position of the paper is very limiting. 

A few basic comments before we get to the main ones.
-1 In the intro:
Please note that when a PDE equation is given, that is complete supervision. Please avoid calling it unsupervised. PDE equation is absolute supervision, sufficient to solve the problem completely. 
-2 In the intro:
Please pay more attention; there is a difference between neural networks and neural operators. One is for learning on the finite-dimensional setting; one is for functional data, that is your case. I only found out on page 8 that the authors talk about Fourier layers, and then I had to go through the referenced paper to realize that the authors use a neural operator for their solver. Even there, please explain you are doing operator learning and using neural operators. Please also define your neural operator architecture. FNO is not even defined or abbreviated in your work. Please bring the abbreviation and explain in your paper that your work is neural operator-based. 

-3 In 3.3 
The authors need to define varrho in the equations in this section. I guessed what that was after reading page 8. There is no varrho in the equations in 3.3.

-4 In 3.3
Why choose linear span? We know such a thing is fine for easy problems and not sufficient for complex problems. That's why even the first paper proposing the PINO paradigm also deals with complex problems and doesn't use linear span simplification. 

-5 In 3.3
 What does Solve even mean in alg2? I only guessed its meaning.

-6 In 3.3. Please be mindful of term solve and term prediction for F. The F predicts something and doesn't solve anything in the sense that solvers in the field of applied math solve. 

-7 In 3.3 
Reading line 304, it feels that the authors first do PINO with only physics and then PINO with data. However, the intro implied that it is the other way. Which is done here?

-8 In 3.3 what is trajectory in line 308?


The main comment is the presentation. Reading the paper, it seems the paper is aiming to advance PINO. However, reading abstract, intro, and other sections up to page 8, the paper does not even talk about PINO. Everything the paper promotes and sugest they are doing is already propsed and establisehd in PINO paper. Please read the PINO paper carefully and possition your paper according to that. Also, in both abstract and intro, mantion what is new in your work. I could not find anuthing different than prior methods. And I am sure there are many things new and novel in this paper, but I just couldn't parse it when abstract and intro are all talking about basics of PINO.

Meta-review: The work proposes to use a neural network to learn a pre-conditioner for the PINN loss depending on some parameter of the PDE problem. Numerics on various problems show that converges speed is significantly increased and a proof of this in the linear setting is provided.

Decision: Accept (Poster)