BiLO: Bilevel Local Operator Learning for PDE inverse problems

## Abstract

Abstract We propose a new neural network based method for solving inverse problems for partial differential equations (PDEs) by formulating the PDE inverse problem as a bilevel optimization problem. At the upper level, we minimize the data loss with respect to the PDE parameters. At the lower level, we train a neural network to locally approximate the PDE solution operator in the neighborhood of a given set of PDE parameters, which enables an accurate approximation of the descent direction for the upper level optimization problem.
The lower level loss function includes the L2 norms of both the residual and its derivative with respect to the PDE parameters.
We apply gradient descent simultaneously on both the upper and lower level optimization problems, leading to an effective and fast algorithm. The method, which we refer to as BiLO (Bilevel Local Operator learning), is also able to efficiently infer unknown functions in the PDEs through the introduction of an auxiliary variable.
Through extensive experiments over multiple PDE systems, we demonstrate that our method enforces strong PDE constraints, is robust to sparse and noisy data,
and eliminates the need to balance the residual and the data loss, which is inherent to the soft PDE constraints in many existing methods.

## 1 Introduction

A fundamental task across various scientific and engineering fields is to infer the unknown parameters of a partial differential equation (PDE) from observed data. Applications include seismic imaging [dengOpenFWILargeScaleMultiStructural2023, martinStochasticNewtonMCMC2012, yangSeismicWavePropagation2021], electrical impedance tomography [uhlmannElectricalImpedanceTomography2009, molinaroNeuralInverseOperators2023], personalized medicine [lipkovaPersonalizedRadiotherapyDesign2019, zhangPersonalizedPredictionsGlioblastoma2024, schaferBayesianPhysicsBasedModeling2021, subramanianEnsembleInversionBrain2023], and climate modeling [senGlobalOptimizationMethods2013].
PDE inverse problems are commonly addressed within the frameworks of PDE-constrained optimization [hinzeOptimizationPDEConstraints2008] or Bayesian inference [stuartInverseProblemsBayesian2010].
In the PDE constrained optimization framework, the objective is to minimize the difference between the observed data and the PDE solution, and the PDE is enforced as a constraint using adjoint, or deep learning, methods.
In the Bayesian inference framework, the inverse problem is formulated as a statistical inference problem, where the goal is to estimate the posterior distribution of the parameters given the data. This requires sampling parameter space and solving the forward PDE multiple times. Here, we develop a constrained optimization framework for solving PDE inverse problems using deep learning.

### 1.1 Related work

The Adjoint Method is a widely used technique for computing the gradients of the objective function with respect to the PDE parameters using numerical PDE solvers in the PDE-constrained optimization framework.
This method provides accurate gradients and strongly satisfies the PDE constraint.
However, the method requires explicitly deriving the adjoint equation and solving the forward and adjoint equations at each iteration, which can be complex and computationally expensive [hinzeOptimizationPDEConstraints2008, plessixReviewAdjointstateMethod2006].

Physics-Informed Neural Networks (PINNs)
have emerged as novel methods for solving inverse problems in a PDE constrained optimization framework [karniadakisPhysicsinformedMachineLearning2021, raissiPhysicsinformedNeuralNetworks2019, jagtapDeepLearningInverse2022, jagtapPhysicsinformedNeuralNetworks2022, chenPhysicsinformedNeuralNetworks2020, zhangPersonalizedPredictionsGlioblastoma2024, yangBPINNsBayesianPhysicsinformed2021, kapoorPhysicsinformedNeuralNetworks2024, chenPhysicsinformedNeuralNetworks2020, jagtapPhysicsinformedNeuralNetworks2022, zhangPersonalizedPredictionsGlioblastoma2024].
PINNs represent PDE solutions using neural networks and embed both the data and the PDE into the loss function through a mesh-free approach.
By minimizing the total loss, PINNs effectively solve the PDE, fit the data, and infer the parameters simultaneously, showcasing integration of mathematical models with data-driven learning processes.
A related approach, Optimizing a Discrete Loss (ODIL), utilizes conventional numerical discretizations of the PDEs and the loss is minimized over the parameters and the PDE solutions at the grid points rather than the weights of a neural network
[karnakovOptimizingDIscreteLoss2022, balcerakIndividualizingGliomaRadiotherapy2024].
However, in these methods, the PDE is enforced as a soft constraint, which requires balancing the residual and the data loss, and can lead to a trade-off between fitting the data and solving the PDE accurately.

Operator Learning aims to train neural networks that approximate the PDE solution operator (parameter-to-solution map) and can serve as surrogate models for the forward PDE solvers [kovachkiNeuralOperatorLearning2022].
Once these surrogates are established, they can be integrated into a Bayesian inference framework or other optimization algorithms to solve inverse problems, leveraging the speed of evaluating a neural network [zhouAIaidedGeometricDesign2023, pathakFourCastNetGlobalDatadriven2022, luMultifidelityDeepNeural2022, maoPPDONetDeepOperator2023].
Some examples of operator learning frameworks include the Fourier Neural Operator [liFourierNeuralOperator2021, liPhysicsInformedNeuralOperator2024, whitePhysicsInformedNeuralOperators2023], DeepONet [luLearningNonlinearOperators2021, wangLearningSolutionOperator2021], In-context operator learning [yangIncontextOperatorLearning2023], among others, e.g. [oleary-roseberryDerivativeInformedNeuralOperator2024, molinaroNeuralInverseOperators2023].
However, for solving the inverse problem, neural operators can encounter challenges when the ground truth is out of the distribution of the training dataset.

There are many other methods for PDE inverse problems using deep learning; see [nganyutanyuDeepLearningMethods2023, herrmannDeepLearningComputational2024, bruntonMachineLearningPartial2023] for more comprehensive reviews.

### Main Contributions

In this work, we focus on solving PDE inverse problems in the PDE-constrained optimization framework using deep learning methods.
The contributions of this paper are as follows:

- •
We formulate the PDE inverse problem as a bilevel optimization problem, where the upper level problem minimizes the data loss with respect to the PDE parameters, and the lower level problem involves training a neural network to approximate the PDE solution operator locally at given PDE parameters, enabling approximation of the descent direction for the upper level optimization problem.
- •
At the lower level problem, we introduce the “residual-gradient” loss, which is the L2 norm of derivative of the residual with respect to the PDE parameters. We show that this loss term compels the neural network to approximate the PDE solution for a small neighborhood of the PDE parameters, thus a “local operator”.
- •
Extensive experiments over multiple PDE systems demonstrate that
our novel formulation is both more accurate and more robust than other
existing methods. It exhibits stronger PDE fidelity, robustness to
sparse and noisy data, and eliminates the need to balance the residual
and the data loss, a common issue in PDE-based soft constraints.
- •
We solve the bilevel optimization problem using gradient descent simultaneously on both the upper and lower level optimization problems, leading to an effective and fast algorithm.
The network architecture is simple and easy to implement.
- •
We extend our method to infer unknown functions that are also parameterized by neural networks through an auxiliary variable. This bypasses the need to learn a high-dimensional local operator.

Our approach combines elements of PINN, operator learning, and the adjoint method.
Our method is closely related to the PINN: both use neural network to represent the solution to the PDE, use automatic differentiation to compute the PDE residual, and aim to solve the PDE and infer the parameters simultaneously.
However, in the PINN, the PDE-constraint is enforced as a regularization term (or soft constraint), leading to a trade-off between fitting the data and solving the PDE accurately, which is the main challenge that we aim to address.
Compared with operator learning, which solves the PDE for a wide range of parameters and requires a large amount of synthetic data for training,
our method only learns the operator local to the PDE parameters at each step of the optimization process and does not require a synthetic dataset for training.
Similar to the adjoint method, we aim to approximate the descent direction for the PDE parameters with respect to the data loss. However, we do not require deriving and solving the adjoint equation.

## 2 Method

### 2.1 PDE Inverse problem as Bi-level optimization

In this section, we present a novel method for solving PDE inverse problems in the framework of PDE-constrained optimization problems using deep learning.
Let $u:\Omega\to\mathbb{R}$ be a function defined over a domain $\Omega\subset\mathbb{R}^{d}$ satisfying some boundary conditions,
and $\hat{u}$ be the observed data, which might be noisy.
Suppose $u$ is governed by a PDE, $F$, which depends on some parameters $\Theta$.
Then the following PDE-constrained optimization problem is solved:

$$ $\displaystyle\min_{\Theta}\quad\lVert u-\hat{u}\lVert^{2}_{2}\quad\quad\textrm {s.t.}\quad F(D^{k}u(\mathbf{x}),...,Du(\mathbf{x}),u(\mathbf{x}),\Theta)=0$ (1) $$

The constraint is a PDE operator that depends on the parameters $\Theta$.
For time-dependent problems, we treat time $t$ as a special component of $\mathbf{x}$, and $\Omega$ includes the temporal domain.

Suppose we know the PDE solution operator (hereafter referred to as the “operator”),
$u(\mathbf{x},\Theta)$, which solves the PDE for any $\Theta$, then we can solve the optimization problem easily by minimizing the objective function using a gradient descent algorithm.
However, finding the full operator $u(\mathbf{x},\Theta)$ is challenging and unnecessary.
Since we are only interested in the descent direction to update $\Theta$, a local approximation of the solution operator suffices, that is,
the operator should approximate the PDE solution for a small neighborhood of a particular value of $\Theta$.
For notational simplicity, we define the residual function of the operator as

$$ $r(\mathbf{x},\Theta):=F(D^{k}u(\mathbf{x},\Theta),...,Du(\mathbf{x},\Theta),u( \mathbf{x},\Theta),\Theta)$ (2) $$

If $u$ is a local operator at $\Theta$, then $r(\mathbf{x},\Theta)=0$ and $\gradient_{\Theta}r(\mathbf{x},\Theta)=0$.
Our goal is to approximate the operator locally at $\Theta$ using a neural network, and then find the optimal PDE parameters $\Theta$ by minimizing the data loss with respect to $\Theta$ using a gradient descent algorithm.

Suppose the local operator is parameterized by a neural network $u(\mathbf{x},\Theta;W)$, where $W$ are the weights of the neural network.
The objective function (1) leads to the following data loss:

$$ $\mathcal{L}_{\rm dat}(\Theta,W)=\frac{1}{|\mathcal{T}_{\rm dat}|}\sum_{\mathbf {x}\in\mathcal{T}_{\rm dat}}\left|u(\mathbf{x},\Theta;W)-\hat{u}(\mathbf{x}) \right|^{2},$ (3) $$

where $\mathcal{T}_{\rm dat}$ is the set of collocation points where the data is observed.
The residual loss is the L2 norm of the residual function

$$ $\mathcal{L}_{\rm res}(W,\Theta):=\frac{1}{|\mathcal{T}_{\rm res}|}\sum_{ \mathbf{x}\in\mathcal{T}_{\rm res}}\left|r(\mathbf{x},\Theta;W)\right|^{2}.$ (4) $$

where $\mathcal{T}_{\rm res}$ is the set of collocation points where the residual loss is evaluated.
We introduce the following loss term, the “residual-gradient loss”, which is the derivative of the residual with respect to the PDE parameters $\Theta$:

$$ $\mathcal{L}_{\rm rgrad}(\Theta,W)=\frac{1}{|\mathcal{T}_{\rm res}|}\sum_{ \mathbf{x}\in\mathcal{T}_{\rm res}}\left|\gradient_{\Theta}r(\mathbf{x},\Theta )\right|^{2},$ (5) $$

Intuitively, this loss compels the neural network to approximate the PDE solution for a small neighborhood of $\Theta$: small variation of $\Theta$ should only lead to small variation of the residual.
If this is satisfied, then the derivative of the data loss with respect to $\Theta$ will approximate the descent direction,
and we can find the optimal $\Theta$ by minimizing the data loss with respect to $\Theta$ using a gradient descent algorithm.
We define the “local operator loss” as the sum of the residual loss and the residual-gradient loss with weight $w_{\rm rgrad}$:

$$ $\mathcal{L}_{\rm LO}(\Theta,W)=\mathcal{L}_{\rm res}(\Theta,W)+w_{\rm rgrad} \mathcal{L}_{\rm rgrad}(\Theta,W)$ (6) $$

Finally, we propose to solve the following bilevel optimization problem:

$$ $\begin{cases}\Theta^{*}=\arg\min_{\Theta}\mathcal{L}_{\rm dat}(\Theta,W^{*}( \Theta))\\ W^{*}(\Theta)=\arg\min_{W}\mathcal{L}_{\rm LO}(\Theta,W)\\ \end{cases}$ (7) $$

In the upper level problem, we find the optimal PDE parameters $\Theta$ by minimizing the data loss with respect to $\Theta$.
In the lower level problem, we train a network to approximate the local operator $u(\mathbf{x},\Theta;W)$ by minimizing the local operator loss with respect to the weights of the neural network.

##### Pre-train and Fine-tune

In this work, we assume access to an initial guess of the PDE parameters, $\Theta_{0}$, alongside their corresponding numerical solution, denoted as $u_{0}$, e.g. from the finite difference method.
The numerical solutions are computed with high accuracy on fine grids, and can be considered as the “exact” solution of the PDE.
We can use the numerical solution to pre-train the neural network, and then use the data to fine-tune the neural network to infer the PDE parameters.
This has been successfully applied in [zhangPersonalizedPredictionsGlioblastoma2024],
and is also similar to curriculum learning, where the neural network learns a “simpler” PDE solution first
[krishnapriyanCharacterizingPossibleFailure2021].
We define the pre-training data loss $\mathcal{L}_{\rm u_{0}}$, which is the MSE between the numerical solution $u_{0}$ and the local operator at $\Theta_{0}$:

$$ $\mathcal{L}_{\rm u_{0}}(W)=\frac{1}{|\mathcal{T}_{\rm res}|}\sum_{\mathbf{x} \in\mathcal{T}_{\rm res}}\left|u(\mathbf{x},\Theta_{0};W)-u_{0}(\mathbf{x}) \right|^{2},$ (8) $$

In the pre-training phase, we solve the following minimization problem

$$ $\min_{W}\mathcal{L}_{\rm LO}(\Theta_{0},W)+\mathcal{L}_{\rm u_{0}}(W)$ (9) $$

The use of $\mathcal{L}_{\rm u_{0}}$ is not mandatory for training the local operator with fixed $\Theta_{0}$, though it can speed up the training process.

### 2.2 Inferring an unknown function

We can also extend our method to learn an unknown function $f(\mathbf{x})$ in the PDE,
such as a variable diffusion coefficient in the Poisson equation or an initial condition in the heat equation.
In these cases, the following PDE constrained optimization problem is solved:

$$ $\displaystyle\min_{f}\quad\lVert u-\hat{u}\lVert^{2}+w_{\rm reg}\lVert \gradient f\rVert^{2}\quad\quad\textrm{s.t.}\quad F(D^{k}u(\mathbf{x}),...,Du( \mathbf{x}),u(\mathbf{x}),f(\mathbf{x}))=0$ (10) $$

where the constraint is a PDE that depends on the unknown function $f$.
Given that these problems are ill-posed, regularization of the unknown function is often necessary.
A typical choice is an L2-norm of the gradient of the unknown function, which penalizes functions that are not smooth.
While the selection of an appropriate regularization form is critical, this paper assumes such choices are predetermined, not an aspect of the method under direct consideration.

Suppose $f$ is parameterized by a neural network $f(\mathbf{x};V)$ with weights $V$.
A straightforward extension from the scalar parameter case is to learn the local operator of the form $u(\mathbf{x},V)$.
However, this would be computationally expensive, as the weights $V$ can be very high dimensional.
We propose to introduce an auxiliary variable $z=f(\mathbf{x})$,
and find a local operator $u(\mathbf{x},z)$ such that $u(\mathbf{x},f(\mathbf{x}))$ solves the PDE locally at $f$.
We define the following function $a$, which is the residual function with an auxiliary variable $z$:

$$ $a(\mathbf{x},z):=F(D^{k}u(\mathbf{x},z),...,Du(\mathbf{x},z),u(\mathbf{x},z),z)$ (11) $$

If $u$ is a local solution operator at $f$, then we should have $a(\mathbf{x},f(\mathbf{x}))=0$ and
$\gradient_{z}a(\mathbf{x},f(\mathbf{x}))=0$.

Suppose $f$ and $u$ are parameterized by neural networks: $f(\mathbf{x};V)$ and $u(\mathbf{x};W)$.
The data loss is similar to the parameter inference case (3) and depends on both $V$ and $W$.
We also need the regularization loss, evaluated on $\mathcal{T}_{\rm reg}$:

$$ $\mathcal{L}_{\rm reg}(V)=\frac{1}{|\mathcal{T}_{\rm reg}|}\sum_{\mathbf{x}\in \mathcal{T}_{\rm reg}}|\gradient_{\mathbf{x}}f(\mathbf{x};V)|^{2}.$ (12) $$

We define the residual loss:

$$ $\mathcal{L}_{\rm res}(W,V):=\frac{1}{|\mathcal{T}_{\rm res}|}\sum_{\mathbf{x} \in\mathcal{T}_{\rm res}}\left|a(\mathbf{x},f(\mathbf{x};V);W)\right|^{2}.$ (13) $$

and the residual-gradient loss:

$$ $\mathcal{L}_{\rm rgrad}(W,V)=\frac{1}{|\mathcal{T}_{\rm res}|}\sum_{\mathbf{x} \in\mathcal{T}_{\rm res}}\left|\nabla_{z}a(\mathbf{x},f(\mathbf{x};V);W)\right |^{2}$ (14) $$

This has the same interpretation as the parameter inference case (5): small variation of $f$ should lead to small variation of the residual.
Finally, we solve the following bilevel optimization problem:

$$ $\displaystyle V^{*}=\arg\min_{V}\mathcal{L}_{\rm dat}(W^{*}(V),V)+w_{\rm reg} \mathcal{L}_{\rm reg}(V)$ (15) $\displaystyle W^{*}(V)=\arg\min_{W}\mathcal{L}_{\rm LO}(W,V)$ (16) $$

where $\mathcal{L}_{\rm LO}=\mathcal{L}_{\rm res}+w_{\rm rgrad}\mathcal{L}_{\rm rgrad}$.
At the upper level, we minimize the data loss and the regularization loss with respect to the weights $V$ of the unknown function, and at the lower level, we minimize the local operator loss with respect to the weights $W$ of the local operator.
The pre-training stage is similar to the parameter inference case.
Given an initial guess of the unknown function $f_{0}$, and its corresponding numerical solution $u_{0}$, we can train the network $f_{V}$ to approximate $f_{0}$ by minimizing the MSE between $f_{V}$ and $f_{0}$,
and train the network $u_{W}$ to be the local operator at $f_{0}$ by minimizing the local operator loss and the MSE between $u_{W}$ and $u_{0}$.

### 2.3 Algorithm

The network architecture involves a simple modification at the input layer (embedding layer) of the typical fully connected neural network: the embedding of the PDE parameters $\Theta$ is randomly initialized and fixed during training, so that the residual-gradient loss can not be made 0 by setting the embedding to 0. See Appendix A for more details.

Solving a bilevel optimization problem is challenging in general [zhangIntroductionBilevelOptimization2023].
In our case, the upper level problem (PDE inverse problem) is usually non-convex, and the lower level problem has a challenging loss landscape [krishnapriyanCharacterizingPossibleFailure2021, basirCriticalInvestigationFailure2022].
However, the lower level problem does not need to be solved to optimality at each iteration because the primary goal is to approximate the descent direction for the upper level problem.
We propose to apply gradient descent to the upper and lower level optimization problems simultaneously.
In Algorithm. 1, we describe our optimization algorithm for inferring scalar parameters in the BiLO framework.
The algorithm for inferring unknown functions is similar.
We write the algorithm as simple gradient descent for notational simplicity while in practice we use the ADAM optimizer [kingmaAdamMethodStochastic2017].

Figure: Algorithm 1 Bi-level Local Operator for inferring scalar PDE parameters

We can have two different learning rates for the two groups of variables $W$ and $\Theta$, denoted as $lr_{W}$ and $lr_{\Theta}$, respectively.
We empirically determined $w_{\rm rgrad}=1e-3$ and $lr_{W}=lr_{\Theta}=0.001$ to be effective across our numerical experiments.
It is not imperative for the residual-gradient loss to be minimized excessively;
it is sufficient that it approximate the correct descent direction. Under somewhat restrictive assumptions, we are able to obtain a theoretical characterization of the bilevel optimization problem (shown below. See Appendix B for a proof). A more general theoretical understanding of the learning dynamics is still lacking and will be left for future work.

Proposition: Assuming
(i) the maximum principal holds for the PDE operator;
(ii) the parametrized local operator $u(W,\Theta)=g$ on $\partial\Omega$ for all $W$ and $\Theta$;
(iii) the lower level problem has a minimizer $W^{*}(\Theta)$ such that the $u(W^{*}(\Theta),\Theta)$ is the local operator,
then the approximate gradient of the upper level objective at $W^{*}(\Theta)$ is exact.

### 2.4 Difference between PINNs and Neural Operators for inverse problems

Neural Operator
Neural operators can serve as surrogate models for PDE solution operators, and can be used in algorithms that require solving the forward PDE multiple times, such as Bayesian inference or derivative-free optimization [kaltenbachSemisupervisedInvertibleNeural2023, luMultifidelityDeepNeural2022], or gradient-based optimization algorithms [zhouAIaidedGeometricDesign2024, luMultifidelityDeepNeural2022, yangRapidSeismicWaveform2023].
However, if the objective is to estimate parameters from limited data, the considerable initial cost for data generation and network training might seem excessive.
The accuracy of specific PDE solutions depends on the accuracy of the neural operator, and
which may decrease if the true PDE parameters fall outside the training data’s distribution [dehoopCostAccuracyTradeOffOperator2022].
This issue can be mitigated by instance-wise fine-tuning using the residual loss [liPhysicsInformedNeuralOperator2024, wangLearningSolutionOperator2021],
though it introduces an additional trade-off: fine-tuning for one parameter set could reduce the operator’s overall accuracy for other parameters.
(an “anchor loss” is thus required to maintain generalization [liPhysicsInformedNeuralOperator2024]).
Thus, in the context of finding the best estimate of the parameters given the data in a PDE-constrained optimization framework, we mainly compare BiLO with PINNs.

PINN
Within the PINN framework, the solution of the PDE is represented by a deep neural network $u(\mathbf{x},W)$, where $W$ denotes all the trainable weights of the neural network [karniadakisPhysicsinformedMachineLearning2021, raissiPhysicsinformedNeuralNetworks2019, luDeepXDEDeepLearning2021].
Notice that the PDE parameters $\Theta$ are not part of the network input. Therefore the data loss does not depend on the PDE parameters $\Theta$, and we write the data loss as $\mathcal{L}_{\rm dat}(W)$.
and the residual loss depends on the weights of the neural network and the PDE parameters.

Solving an inverse problem using PINN involves minimizing an unconstrained optimization problem, where the objective function is the weighted sum of the residual loss and the data loss

$$ $\min_{W,\Theta}\mathcal{L}_{\rm res}(W,\Theta)+w_{\rm dat}\mathcal{L}_{\rm dat }(W)$ (19) $$

where $w_{\rm dat}$ is the weight of the data loss.
For simplicity of discussion, we assume the weight of the residual loss is always 1.
The key feature is that the PDE is enforced as a soft constraint, or as a regularization term for fitting the data.
The relationship between the PDE parameter and the data loss is indirect:
the descent directions of the PDE parameters are given by $\gradient_{\Theta}\mathcal{L}_{\rm res}$, which are independent of the data loss.

Challenges for PINNs
Solving PDE inverse problems using PINNs can encounter challenges stemming from the soft PDE constraint (19), especially when the data is sparse and noisy, or when the PDE model does not fully explain the data.
The soft PDE constraint can result in a trade-off between fitting the data and solving the PDE accurately.
In addition, since the PDE parameters are updated in the descent direction of the residual loss, they can be biased toward parameters corresponding to very smooth solutions.
This can also lead to slow convergence, as a small residual loss can be achieved by any PDE parameters corresponding to a smooth solution, and the step size can become small before the data loss is sufficiently minimized.
It is important to recognize that PINNs can indeed be effective for PDE inverse problems, particularly when the data is abundant and the noise is independent and identically distributed:
in this case, the minimizer of the data loss still gives a good approximation of the PDE solution [kapoorPhysicsinformedNeuralNetworks2024, chenPhysicsinformedNeuralNetworks2020, jagtapPhysicsinformedNeuralNetworks2022].
The challenges become more pronounced with sparse and noisy data or misspecified models [zhangPersonalizedPredictionsGlioblastoma2024].

There are many techniques to improve the performance of PINNs,
such as adaptive sampling and weighting of collocation points [nabianEfficientTrainingPhysicsinformed2021, wuComprehensiveStudyNonadaptive2023, luDeepXDEDeepLearning2021, anagnostopoulosResidualbasedAttentionPhysicsinformed2024], new architectures [jagtapAdaptiveActivationFunctions2020, wangPirateNetsPhysicsinformedDeep2024, wangUnderstandingMitigatingGradient2021, moseleyFiniteBasisPhysicsinformed2023],
new optimization algorithms [basirPhysicsEqualityConstrained2022, krishnapriyanCharacterizingPossibleFailure2021], new loss functions [wangPhysicsInformedLoss2022, yuGradientenhancedPhysicsinformedNeural2022, sonSobolevTrainingPhysics2021], adaptive weighting of loss terms [madduInverseDirichletWeighting2022, wangUnderstandingMitigatingGradient2021, mcclennySelfAdaptivePhysicsInformedNeural2022, wangExpertGuideTraining2023].
While most of these techniques focus on the forward problem, they can potentially be used for inverse problem as well.
However, these techniques do not fundamentally change the soft PDE-constraints in the PINN framework.
In our work, we propose a different optimization problem that does not involve a trade-off between the residual loss and the data loss,
and our method can be used in conjunction with many of these techniques to improve the performance.
Therefore, in the following numerical experiments, we do not use any of these techniques, and we focus on comparing the two different optimization formulations (bi-level optimization and the soft PDE-constraints).

## 3 Numerical Experiments

In Section 3.1, we infer two scalar parameters in the Fisher-KPP equation and compare the performance of BiLO, PINN and DeepONet.
In Section 3.2, we infer an unknown function in the Poisson equation and compare the performance of BiLO and PINN (results of DeepONet are shown in Appendix D.2).
We denote the neural network solution (from BiLO, PINN, or DeepONet) by $u_{\rm NN}$, and denote the numerical solution with the inferred parameters using the Finite Difference Method (FDM) by $u_{\rm FDM}$, which is solved to a high accuracy.
A large discrepancy between $u_{\rm NN}$ and $u_{\rm FDM}$ suggests that the PDE is not solved accurately by the neural network.

We provides the training detail and hyperparameters for the numerical experiments in Section 3.1 and 3.2 in Appendix C.
Details of the DeepONet architecture and training are provided in Appendix D.
Appendix E provides additional numerical experiments:
(1) E.1 Inferring a scalar coefficient in a nonlinear ODE problem;
(2) E.2 Inferring the initial condition of a 1D heat equation;
(3) E.3 Inferring the variable diffusion coefficient of a 2D Poisson problem, where we achieve better or comparable performance as in PINO [liPhysicsInformedNeuralOperator2024].

### 3.1 Fisher-KPP Equation

In this example, we aim to infer the unknown parameters $D$ and $\rho$ in the following Fisher-KPP equation [zouCorrectingModelMisspecification2024], which is a nonlinear reaction-diffusion equation:

$$ $\begin{cases}u_{t}(x,t)=0.01Du_{xx}(x,t)+\rho u(1-u)\\ u(x,0)=\frac{1}{2}\sin(\pi x)^{2}\\ u(0,t)=u(1,t)=0\\ \end{cases}$ (20) $$

The initial guesses of the PDE parameters are $D_{0}=1$ and $\rho_{0}=1$, and the ground truth parameters are $D_{GT}=2$ and $\rho_{GT}=2$.
This equation has been used to model various biological phenomena, such as
the growth of tumors [swansonQuantitativeModelDifferential2000, harpoldEvolutionMathematicalModeling2007]
or the spreading of misfolded proteins [zhangDiscoveringReactionDiffusion2024, schaferNetworkDiffusionModeling2020, schaferBayesianPhysicsBasedModeling2021].
In our tests the data is only provided at the final time $t=1$, which is more challenging than the case where data is provided at multiple time points.
This single-time inference problem has application in patient-specific parameter estimation of tumor growth models using medical images, where only one time point may be available, e.g., in the case of glioblastoma [balcerakIndividualizingGliomaRadiotherapy2024, zhangPersonalizedPredictionsGlioblastoma2024, ezhovLearnMorphInferNewWay2023, scheufeleFullyAutomaticCalibration2021].

Effect of residual-gradient loss
We plot the trained local operator $u(\mathbf{x},D_{0}+\delta D,\rho_{0}+\delta\rho;W)$ at $t=1$,
for $(\delta D,\delta\rho)$ = (0.5,0) and (0,0.1),
and the corresponding FDM solution in Fig. 1 (a).
We can see that even though the network is only trained using the initial parameters, because of the residual-gradient loss, the network can approximate the solution of the PDE for a small neighborhood of the parameters.
This suggests that the derivative of the data loss with respect to the parameters should give the correct descent direction.

Trajectory of the Parameters
We consider the case without noise and show the trajectories of the parameters $D$ and $\rho$ during the fine-tuning process in Fig. 1 (b).
Each BiLO trajectory (black line) corresponds to a different random initialization of the neural network, and are obtained by our simultaneous gradient descent. They roughly follow the trajectory that is obtained by solving the lower level problem to a small tolerance before updating the PDE parameters (red dashed line).
The contours are the data loss in log scale using the FDM solution for each parameter pair $(D,\rho)$.
Note that the contour lines do not represent the actual loss landscape of our optimization problem, since at each step we are not solving the PDE to high accuracy.
From the landscape we can also see that single-time inference is challenging, as the gradient with respect to $D$ is much smaller than $\rho$, leading to a narrow valley in the loss landscape along the $D$-direction.

Figure: Figure 1: (a) Visualization of the local operator $u(\mathbf{x},D_{0}+\delta D,\rho_{0}+\delta\rho;W)$ at $t=1$ for $\delta D=0.5$ or $\delta\rho=0.2$, and the corresponding FDM solutions. (b) Trajectory of the parameters $D$ and $\rho$ during fine-tuning roughly follow the path of the steepest descent. The dashed line is the trajectory when the lower level problem is solved to a small tolerance. The contours correspond to the data loss in log scale, computed using the FDM solution.
Refer to caption: extracted/5794873/fig_fk_var.png

Inference with noise
In this experiment, we consider inference under noise $\epsilon\sim N(0,1e-4)$.
In Fig. 2, we show the results of BiLO and PINNs with different weights $w_{\rm dat}$= 0.01, 0.1, 1.
We can see that for $w_{\rm dat}$ = 0.01 and 0.1, the PDE is solved relatively accurately, since $u_{\rm NN}$ and $u_{\rm FDM}$ overlap.
For $w_{\rm dat}$ = 1, the PDE is not solved accurately and the network is over-fitting the data.
In addition, PINNs have difficulties in obtaining accurate estimates of $D$ due to the challenging loss landscape.
Our new method gives more accurate inferred parameter and PDE solution.

Figure: Figure 2: Enlarged view of the network predicted solutions $u_{\rm NN}$ (BiLO and PINNs with different $w_{\rm dat}$) and FDM solutions $u_{\rm FDM}$ at final time, in the region $(x,u)\in[0.2,0.8]\times[0.6,0.85]$. BiLO gives more accurate inferred parameters and PDE solution.
Refer to caption: extracted/5794873/fig_fk_data.png

In Table 1, we show the mean and standard deviation (std) of various metrics for BiLO, PINNs with different $w_{\rm dat}$, and DeepONets with different pretraining datasets.
The ground truth solution should have an average data loss of $\mathcal{L}_{\rm dat}=1e-4$, which is the variance of the noise.
We can see that the loss landscape is particularly challenging, leading to relatively large error in $D$ for all methods.
For the PINN, we see that $w_{\rm dat}=0.01$ leads to under-fitting of the data, as the data loss is larger than the variance of the noise;
and $w_{\rm dat}=$ 10 shows clear sign of over-fitting of the data, as the data loss is getting smaller than the variance of the noise.
The DeepONets are first pretrained with numercial solutions of the PDE with various $D$ and $\rho$.
Then a gradient-based optimization algorithm is used to solve the inverse problem.
We consider both coarse and dense sampling of the parameters $D$ and $\rho$ that include the ground truth parameters. Additionally, we also consider a dense sampling but the ground truth parameters are out-of-distribution(OOD). Details are provided in Appendix. D.1.
We can see that the results from DeepONet are affected by the quality of the pretraining dataset.
Overall, BiLO gives more accurate inferred parameters and PDE solution, is robust to the noise, and does not require a large amount of pretraining data.

**Table 1: Comparison of BiLO, PINNs (with various $w_{\rm dat}$) and DeepONet (with various pretraining dataset) for a Fisher-KPP PDE problem with noise $\epsilon\sim N(0,1e-4)$. BiLO gives more accurate inferred parameters and PDE solution.**
| method | $\lvert D-D_{GT}\lvert$ | $\lvert\rho-\rho_{GT}\lvert$ | $\lVert u_{\rm NN}-u_{\rm FDM}\lVert_{\infty}$ | $\mathcal{L}_{\rm data}$ |
| --- | --- | --- | --- | --- |
| BiLO | 0.26$\pm$0.10 | 0.06$\pm$0.03 | 3.36e-3$\pm$1.14e-3 | 1.01e-4$\pm$2.77e-5 |
| PINN(1e-1) | 0.85$\pm$0.07 | 0.17$\pm$0.02 | 9.40e-3$\pm$9.15e-4 | 1.43e-4$\pm$2.58e-5 |
| PINN(1e0) | 0.40$\pm$0.13 | 0.09$\pm$0.03 | 4.41e-3$\pm$1.44e-3 | 8.68e-5$\pm$3.00e-5 |
| PINN(1e1) | 0.44$\pm$0.21 | 0.10$\pm$0.04 | 4.93e-3$\pm$2.10e-3 | 3.29e-5$\pm$2.02e-5 |
| DeepONet(Coarse) | 0.95$\pm$0.74 | 0.24$\pm$0.20 | 7.96e-3$\pm$6.36e-3 | 6.26e-5$\pm$2.81e-5 |
| DeepONet(Dense) | 0.48$\pm$0.40 | 0.13$\pm$0.10 | 4.85e-3$\pm$3.47e-3 | 6.23e-5$\pm$1.95e-5 |
| DeepONet(OOD) | 0.95$\pm$0.86 | 0.35$\pm$0.38 | 1.62e-2$\pm$1.75e-2 | 6.18e-5$\pm$1.88e-5 |

### 3.2 Poisson Equation with Variable Diffusion Coefficient

In this test, we consider the following Poisson equation on $[0,1]$ with $u(0)=u(1)=0$:

$$ $(D(x)u^{\prime}(x))^{\prime}=-\pi^{2}\sin(\pi x)$ (21) $$

and aim to infer the variable diffusion coefficient $D(x)$ such that $D(0)=D(1)=1$.
The ground truth $D(x)$ is a “hat” function $D(x)=1+2x$ for $x\in[0,0.5)$ and $D(x)=2-2x$ for $x\in[0.5,1]$.
We start with initial guess $D_{0}(x)=1$.

Effect of residual-gradient loss
In Fig. 3, we visualize the local operator $u(x,z;W)$ after pre-training with $D_{0}(x)=1$.
We consider the variation
$\delta D_{1}(x)=-0.1$, and
$\delta D_{2}(x)=0.1x$
and evaluate the neural network at $u(x,D_{0}(x)+\delta D_{i}(x);W)$ for $i=1,2$.
The FDM solutions of the PDE corresponding to $D_{0}(x)+\delta D_{i}(x)$ are also plotted.
We can see that the neural network approximates the solution corresponding to $D_{0}(x)+\delta D_{i}(x)$ well.

Figure: Figure 3: (a) Visualizing the operator $u(x,D(x)+\delta D(x);W)$ after pre-training with $D_{0}(x)=1$. (b) Inferred $D(x)$ using noise data with 5 random seeds: BiLO and PINN with various $w_{\rm dat}$. BiLO gives more accurate inferred $D(x)$.
Refer to caption: extracted/5794873/fig_varpoi_var.png

**Table 2: Comparison of BiLO and PINNs with various $w_{\rm dat}$ for inferring a variable diffusion coefficient from noisy data. BiLO is more robust to the noise and gives a more accurate inferred diffusion coefficient and PDE solution.**
| method | $\lVert D-D_{GT}\lVert_{\infty}$ | $\lVert D-D_{GT}\lVert_{2}$ | $\lVert u_{\rm NN}-u_{\rm FDM}\lVert_{\infty}$ | $\mathcal{L}_{\rm data}$ |
| --- | --- | --- | --- | --- |
| BiLO | 5.86e-2$\pm$1.99e-2 | 2.01e-2$\pm$7.94e-3 | 3.94e-3$\pm$1.93e-3 | 1.01e-4$\pm$1.80e-5 |
| PINN(1e0) | 9.99e-2$\pm$2.88e-3 | 3.97e-2$\pm$1.73e-3 | 6.37e-3$\pm$1.54e-3 | 1.09e-4$\pm$1.85e-5 |
| PINN(1e1) | 8.61e-2$\pm$7.50e-3 | 3.25e-2$\pm$3.96e-3 | 4.43e-3$\pm$1.37e-3 | 1.02e-4$\pm$1.87e-5 |
| PINN(1e2) | 7.13e-2$\pm$1.59e-2 | 3.11e-2$\pm$1.09e-2 | 4.88e-3$\pm$1.45e-3 | 9.42e-5$\pm$1.55e-5 |

Inference With Noise Data
In this experiment, we consider inference under noise $\epsilon\sim N(0,1\text{e-}4)$ and set $w_{\rm reg}=$1e-3.
In Table. 2, we show the mean and standard deviation of various metrics.
We consider the $L_{\infty}$ and $L_{2}$ error of the inferred function $D(x)$ from the ground truth $D_{GT}$, which measure the accuracy of the inferred function;
The $L_{\infty}$ error between $u_{\rm NN}$ and $u_{\rm FDM}$ indicates the accuracy of neural network solution;
We also show the average data loss $\mathcal{L}_{\rm dat}$, which ideally should be close to the variance of the noise (1e-4).
A smaller or larger average $\mathcal{L}_{\rm dat}$ indicates tendencies to under-fitting or over-fitting of the data.
For the PINNs, we can see that the optimal $w_{\rm dat}$ is about 10, as increasing to 100 leads to over-fitting of the data, and decreasing to 0.1 leads to under-fitting of the data.
BiLO results in more accurate inferred diffusion coefficient and PDE solution, and is robust to the noise.
We also visualize the inferred $D(x)$ in Fig. 3 (b).
For the PINN, we can see that a small $w_{\rm dat}$ leads to very smooth $D(x)$, while a large $w_{\rm dat}$ leads to an oscillating $D(x)$ due to over-fitting.
BiLO gives more accurate inferred $D(x)$ that better approximate the kink of the ground truth $D(x)$.
In Appendix. D.2, we also compare BILO with DeepONet, whose performance depends on the pretraining dataset.

## 4 Conclusion

In this work, we propose a Bi-level Local Operator (BiLO) learning framework for solving PDE inverse problems:
we minimize the data loss with respect to the PDE parameters at the upper level, and learn the local solution operator of the PDE at the lower level.
The bi-level optimization problem is solved using simultaneous gradient descent, leading to an efficient algorithm.
Empirical results demonstrate more accurate parameter recovery and stronger fidelity to the underlying PDEs under sparse and noisy data,
compared with the soft PDE-constraint formulation, which faces the delicate trade-off between adhering to the PDE constraints and accurately fitting the data.
As limitations: (1) the convergence results are mainly empirical with limited theoretical analysis, (2) the numerical experiments are limited to low dimensional problems, and (3) the architecture of the neural network is simple.
Future work includes theoretical analysis of the method, applying the method to more complex and higher dimensional problems, and improving the network architectures.

## Code Availability

The code for the numerical experiments is available at [https://github.com/Rayzhangzirui/BILO](https://github.com/Rayzhangzirui/BILO).

## Acknowledgment

R.Z.Z and J.S.L thank Babak Shahbaba for the GPU resources. J.S.L acknowledges partial support from the National Science Foundation through grants DMS-2309800, DMS-1953410 and DMS-1763272 and the Simons Foundation (594598QN) for an NSF-Simons Center for Multiscale Cell Fate Research.

Supplementary Material

In Appendix A, we provide the network architecture of the neural network used in the numerical experiments.
In Appendix B, we provide a justification of the simultaneous gradient descent algorithm for the bi-level optimization problem.
Appendix C provides the training detail and hyperparameters for the numerical experiments in Section 3.1 and 3.2 in the main text.
In Appendix D, we compare BILO with solving PDE inverse problems using a neural operator.
Appendix E includes numerical experiments, which includes

- •
E.1 Inferring a scalar coefficient in a nonlinear ODE problem
- •
E.2 Inferring the initial condition of a 1D heat equation
- •
E.3 Inferring the variable diffusion coefficient of a 2D Poisson problem

Appendix F shows the computational cost of BiLO.

## Appendix A Network Architecture

The network architecture involves a simple modification at the input layer (embedding layer) of the typical fully connected neural network.
For the scalar parameter case, the input layer maps the inputs $\mathbf{x}$ and the unknown PDE parameters $\Theta$ to a high-dimensional vector $\mathbf{y}$, using an affine transformation followed by a non-linear activation function $\sigma$:

$$ $\mathbf{y}=\sigma(W\mathbf{x}+R\Theta+\mathbf{b}),$ (22) $$

where $W$ is the embedding matrix for $\mathbf{x}$, $R$ is the embedding matrix for $\Theta$, and $\mathbf{b}$ is the bias vector.
The key is that the embedding matrix $R$ should be non-trainable.
Otherwise, $\mathcal{L}_{\rm rgrad}(W,\Theta)$ can be made 0 by setting $R$ to be 0.
In our work, $R$ will be randomly initialized in the same way as $W$, using uniform distributions in the range of $[-1/\sqrt{d},1/\sqrt{d}]$, where $d$ is the number of input units in the layer.
The embedding vector $\mathbf{y}$ is then passed through a series of fully connected layers with activation functions.
The output of the network is denoted as $\mathcal{N}(\mathbf{x},\Theta;W)$, where $W$ denotes all the trainable weights of the neural network.
In some cases, a final transformation is applied to the output of the neural network
$u(\mathbf{x};W)=\tau\left(\mathcal{N}(\mathbf{x},\Theta;W),\mathbf{x}\right)$,
to enforce the boundary condition [dongMethodRepresentingPeriodic2021, luPhysicsInformedNeuralNetworks2021, sukumarExactImpositionBoundary2022].

## Appendix B Simultaneous Gradient Descent

In the main text, we describe the simultaneous gradient descent algorithm for the bi-level optimization problem.
In this section, we provide a justification of the algorithm under some assumptions.

We consider the boundary value problem:

$$ $\begin{cases}\mathbf{L}u=f&\text{in }\Omega\\ u=g&\text{on }\partial\Omega,\end{cases}$ (23) $$

where $\Omega$ is an connected, open and bounded subset of $\mathbb{R}^{d}$. $\mathbf{L}$ denoteds a second-order parital differential operator:

$$ $\mathbf{L}u=\sum_{i,j=1}^{d}a_{ij}\partial_{ij}u+\sum_{i=1}^{d}b_{i}\partial_{ i}u+cu$ (24) $$

where the coefficients $a_{ij}$, $b_{i}$, $c$ are colletively denoted as $\Theta$.
We denote $\mathbf{L}_{\Theta}$ as the derivative of $\mathbf{L}$ with respect to $\Theta$, which is also a differential operator.

We say a function $u(\mathbf{x},\Theta)$ is a local solution operator of the PDE (24)
at $\Theta$ if (1) $\mathbf{L}u=f$ and (2) $L_{\Theta}u+L\nabla_{\Theta}u=0$. That is, the residual at $\Theta$ is zero and the gradient of the residual w.r.t $\Theta$ is zero.

We consider a parameterized local operator $u(\mathbf{x},\Theta;W)$.
For notational simplicity, we omit the dependence of $u$ on $\mathbf{x}$ in the following discussion.
We assume that $u(\Theta;W)=g$ on $\partial\Omega$ for all $W$ and $\Theta$.

Our bilevel optimizaiton problem is

$$ $\min_{\Theta}\int_{\Omega}\left(u(\Theta,W^{*}(\Theta))-\hat{u}\right)^{2}d \mathbf{x}$ $$

$$ $W^{*}(\Theta)=\arg\min\int_{\Omega}\left(\mathbf{L}u-f\right)^{2}+w_{\rm reg} \left(L_{\Theta}u+L\nabla_{\Theta}u\right)^{2}d\mathbf{x}$ $$

where $\mathbf{L}u-f$ is the residual of the PDE, and $L_{\Theta}u+L\nabla_{\Theta}u$ is the gradient of the residual w.r.t $\Theta$.

In our simultaneous gradient descent, the gradient of the upper level objective with respect to $\Theta$ is given by

$$ $g_{\rm a}(W,\Theta)=\int_{\Omega}\left(u(W,\Theta)-\hat{u}\right)\left(\nabla_ {\Theta}u(W,\Theta)\right)d\mathbf{x}$ (25) $$

The exact gradient of the upper level objective is

$$ $g(\Theta)=\int_{\Omega}\left(u(W^{*}(\Theta),\Theta)-\hat{u}\right)\left( \nabla_{W}u(W^{*}(\Theta),\Theta)\nabla_{\Theta}W^{*}(\Theta)+\nabla_{\Theta}u (W^{*}(\Theta),\Theta)\right)d\mathbf{x}$ (26) $$

At $W^{*}(\Theta)$, the difference between the exact gradient and the approximate gradient, which we denote as $\Delta g$, is given by

$$ $\displaystyle\Delta g(\Theta):$ $\displaystyle=g_{\rm a}(W^{*}(\Theta),\Theta)-g(\Theta)$ (27) $\displaystyle=\int_{\Omega}\left(u(W^{*}(\Theta),\Theta)-\hat{u}\right)\left( \nabla_{W}u(W^{*}(\Theta),\Theta)\nabla_{\Theta}W^{*}(\Theta)\right)d\mathbf{x}$ $$

Suppose the lower level problem has a minimizer $W^{*}(\Theta)$ such that the $u(W^{*}(\Theta),\Theta)$ is the local operator.

$$ $\mathbf{L}u(W^{*}(\Theta),\Theta)-f=0$ (28) $$

and

$$ $\mathbf{L}_{\Theta}u(W^{*}(\Theta),\Theta)+\mathbf{L}\nabla_{\Theta}u(W^{*}( \Theta),\Theta)=0$ (29) $$

Take the derivative of the Eq. (28) with respect to $\Theta$, we have

$$ $\mathbf{L}_{\Theta}u(W^{*}(\Theta),\Theta)+\mathbf{L}\nabla_{\Theta}u(W^{*}( \Theta),\Theta)+\mathbf{L}\nabla_{W}u(W^{*}(\Theta),\Theta)\nabla_{\Theta}W^{* }(\Theta)=0$ (30) $$

From Eq. (29) and Eq. (30), we have

$$ $\mathbf{L}\nabla_{W}u(W^{*}(\Theta),\Theta)\nabla_{\Theta}W^{*}(\Theta)=0$ (31) $$

We denote the function
$v:=\nabla_{W}u(W^{*}(\Theta),\Theta)\nabla_{\Theta}W^{*}(\Theta)$.
Since $u(W,\Theta)=g$ on $\partial\Omega$ for all $W$ and $\Theta$,
we have $v=0$ on $\partial\Omega$.
Therefore, we have $\mathbf{L}v=0$ in $\Omega$ and $v=0$ on $\partial\Omega$.
If the maximum principal holds for the operator $\mathbf{L}$, for example, when $\mathbf{L}$ uniformly elliptic and $c\geq=0$, [evansPartialDifferentialEquations2010]
then we have $v=0$.

By Cauchy-Schwarz inequality, we have

$$ $||\Delta g||_{2}\leq||u(W^{*}(\Theta),\Theta)-\hat{u}||_{2}||v||_{2}=0$ (32) $$

That is, the approximate gradient at $W^{*}(\Theta)$ is exact.

We summarize the above discussion in the following proposition:

Proposition: Assuming
(i) the maximum principal holds for $\mathbf{L}$;
(ii) the parametrized local operator $u(W,\Theta)=g$ on $\partial\Omega$ for all $W$ and $\Theta$;
(iii) the lower level problem has a minimizer $W^{*}(\Theta)$ such that the $u(W^{*}(\Theta),\Theta)$ is the local operator,
then the approximate gradient (25) of the upper level objective at $W^{*}(\Theta)$ is exact.

The assumptions are more restrictive than the numerical experiments.
For example, in the Fisher-KPP example, the PDE operator is nonlinear.
A more comprehensive and general analysis is left for future work.

## Appendix C Training Details

For each numerical experiment, we solve the optimization problem 5 times with different random seed.
which affect both the initialization of the neural network and the noise in the data (if applicable).
Although each realization of the noise may yield a different optimal parameter $\Theta^{*}$, the average of the optimal parameters across multiple runs should still be close to the ground truth parameter $\Theta_{GT}$.
Therefore, we report the mean and standard deviation of the error between the inferred parameters, or functions, and the ground truth quantities.

In all the numerical experiment, we use the tanh activation function and 2 hidden layers, each with 128 neurons, for both PINN and BiLO.
The collocation points are evenly spaced as a grid in the domain.
For all the optimization problems, we use the Adam optimizer with learning rate 0.001 and run a fixed number of steps.

##### Fisher-KPP Equation

Our local operator take the form of $u(x,t,D,\rho;W)=u(x,0)+\mathcal{N}(x,t,D,\rho;W)x(1-x)t$ so that the initial condition and the boundary condition are satisfied.
Let $X_{r}$, $X_{d}$ be the spatial coordinates evenly spaced in $[0,1]$, and $T_{r}$ be temporal coordinates evenly spaced in $[0,1]$.
We set $\mathcal{T}_{\rm res}=X_{r}\times T_{r}$ and $|X_{r}|=|T_{r}|=21$, that is, the residual collocation points are a uniform grid in space and time.
We set $\mathcal{T}_{\rm dat}=X_{d}\times\{1\}$ and $|X_{d}|=11$, that is, the data collocation points form a uniform grid at the final time $t=1$.
Both BiLO and PINN are pretrained with the initial guess for 10,000 steps, and fine-tuned for 10,000 steps.

In Fig. 4, we show an example of the training history of the losses and the inferred parameters corresponding to Fig 2, and indicate the ground truth with grey dashed line.
For data loss, the ground truth solution should have a data loss of 1e-4, which is the variance of the noise.

Figure: Figure 4: Traning history of losses and the inferred parameters corresponding to Fig 2. The losses are in log scale.
Refer to caption: extracted/5794873/fig_fk_loss.png

##### Poisson Equation with Variable Diffusion Coefficient

The local operator takes the form of $u(x,z;W)=\mathcal{N}_{1}(x,z;W)x(1-x)$ to enforce the boundary condition, where the fully connected neural network $\mathcal{N}_{1}$ has 2 hidden layers, each with 128 neurons.
The unknown function is parameterized by $D(x;V)=\mathcal{N}_{2}(x,V)x(1-x)+1$, where $\mathcal{N}_{2}$ has 2 hidden layers, each with 64 neurons.
For pre-training, we set $|\mathcal{T}_{\rm res}|=|\mathcal{T}_{\rm reg}|=|\mathcal{T}_{\rm dat}|=101$, and train 10,000 steps.
For fine-tuning, we set $|\mathcal{T}_{\rm res}|=|\mathcal{T}_{\rm reg}|=101$ and $|\mathcal{T}_{\rm dat}|=51$, and train 10,000 steps.

## Appendix D Comparison with Neural Operators

In this section, we compare the results of BiLO and Neural Operators (NO) for solving the inverse problems.
For the NO, we use the DeepONet architecture [luLearningNonlinearOperators2021] as an example, which is shown to have comparable performance with FNO [liFourierNeuralOperator2021, luComprehensiveFairComparison2022].

It is difficult to directly compare the performance of NO and PINN/BiLO, since NOs are designed to learn the solution operator of the PDE, while both the PINN and BiLO can be considered as the solver of the PDE, which solve the PDE for one set of parameters.
Ususally, NO is trained with a large amount of numerical solutions.
In this experiment, for solving the inverse problem, we first train the NO, and then we use the NO as a surrogate and use gradient-based optimization to infer the parameters of the PDE.
We show the the quality of the inferred parameters depends on the quality of the synthetic data used to train the NO.

We emphasize that NO can excel in multi-query scenarios, such as sovling the inverse problem in a Bayesian framework, which requires evaluating the solution of the PDE for many different parameters.

### D.1 Fisher-KPP Equation

In this experiment, we consider the Fisher-KPP equation with noise, as in Section. 3.1.
We consider the following 3 datasets for pretraining the DeepONet.
The ground truth parameters are $D_{GT}=2$ and $\rho_{GT}=2$, and the initial guess is $D_{0}=1$ and $\rho_{0}=1$.
The PDE pararameters are sampled with different range and different resolution. We use the notation $a:h:b$ to denote an array from $a$ to $b$ with step $h$.

- •
Coarse: $D=0.8:0.05:3$, $\rho=0.8:0.05:3$.
- •
Dense: $D=0.8:0.02:3$, $\rho=0.8:0.02:3$.
- •
Out-of-distribution (OOD): $D=0.8:0.02:1.8$, $\rho=0.8:0.02:1.8$.

In the “Coarse” dataset, the parameters are sampled with a larger step size. In the “Dense” dataset, the parameters are sampled with a smaller step size. In the “OOD” dataset, the parameters are sampled with a smaller step size, does not include the ground truth parameters.

We use the following architecture for the DeepONet:

$$ $G_{W}(D,\rho,\mathbf{x})=\sum_{i=1}^{k}b_{k}(D,\rho)t_{k}(\mathbf{x})$ $$

where $b_{k}(D,\rho)$ is the k-th output of the “branch net”, and $t_{k}(\mathbf{x})$ is the k-th output of the “truck net”. Both the trunk net and the truck net are parameterized by fully neural networks with 2 hidden layers, each with 128 neurons, so that the total number of parameters (46179) are comparable to the network used by BILO (42051).
The weights of the DeepONet are denoted as $W$.
A final transformation on the output $G_{W}$ is used to enforce the boundary condition.
We pre-train multiple DeepONets with 10,000 steps using each datasets.

Given a pretrain dataset with collections of {$D^{j}$, $\rho^{j}$}
and their corresponding solutions $u^{{j}}$ for $j=1,\ldots,m$,
we first train the DeepONet with the following operator data loss:

$$ $\min_{W}\sum_{j=1}^{m}\sum_{\mathbf{x}\in\mathcal{T}_{\rm dat}}\left|G_{W}(D^{ j},\rho^{j},\mathbf{x})-u^{j}(\mathbf{x})\right|^{2}$ $$

where $\mathcal{T}_{\rm dat}$ is the same as those used in the BiLO and PINN.
For the inverse problem, we fix the weights $W$ and treat the $D$ and $\rho$ as unknown variables.
We minimize the data loss:

$$ $\min_{D,\rho}\frac{1}{|\mathcal{T}_{\rm dat}|}\sum_{\mathbf{x}\in\mathcal{T}_{ \rm dat}}\left|G_{W}(D,\rho,\mathbf{x})-\hat{u}(\mathbf{x})\right|^{2}$ $$

where $\hat{u}$ is the noisy data.

As shown in Table 1 in the main text, the performance of the inference depends on properties of the pre-training dataset.
When the ground truth is out of the distribution of the pre-training dataset, the DeepONet gives poor performance.

### D.2 Variable-Diffusion Coefficient Poisson Equation

In this experiment, we infer the variable diffusion coefficient $D(x)$ in the Poisson equation using a DeepONet.
The pretrain dataset is generated by solving the Poisson equation with 1000 samples of variable diffusion coefficient $D(x)$.
$D(x)$ is sampled from a Gaussina Random field on $[0,1]$, conditioned on $D(0)=D(1)=1$.
The covariance function is the gaussian kernel, with variance 0.05 and different length scale $l=0.2,0.3,0.4$.
See Figure 5 for the samples of $D(x)$ and their corresponding solutions.
As $l$ increases, the samples of $D(x)$ become smoother.

Figure: Figure 5: Samples (gray lines) of $D(x)$ with various length scale $l$ and their corresponding solutions. Black line is the ground truth $D$ and $u$
Refer to caption: extracted/5794873/fig_op_varpoi_sample.png

The DeepONet has the following architecture:

$$ $G_{W}(\mathbf{D},\mathbf{x})=\sum_{i=1}^{k}b_{k}(\mathbf{D})t_{k}(\mathbf{x})$ $$

where the vector $\mathbf{D}$ respresent the values of $D(x)$ at the collocation points. A final transformation on the output $G_{W}$ is used to enforce the boundary condition.
In this experiment, both $D$ and $u$ are evaluated at 101 points in $[0,1]$.
Let $x_{i}$ be the collocation points in $[0,1]$ for $i=1,\ldots,N$. Let $\{D^{j}(x_{i}),u^{j}(x_{i})\}$ be the samples of $D$ and the corresponding solutions $u$ at $x_{i}$ for $j=1,\ldots,m$. We denote $\mathbf{D}^{j}$ as the vector of $D^{j}(x_{i})$ for $i=1,\ldots,N$.
In the pre-training step, we solve the following minimization problem

$$ $\min_{W}\sum_{j=1}^{m}\sum_{i=1}^{N}\left|G_{W}(\mathbf{D}^{j},x_{i})-u^{j}(x_ {i})\right|^{2}$ $$

For the inverse problem, we fix the weights $W$ and treat the $\mathbf{D}$ as an unknown variable.
We minimize the data loss and a finite difference discretizations of the regularization term $|D(x)|^{2}$:

$$ $\min_{\mathbf{D}}\frac{1}{N}\sum_{i=1}^{N}\left|G_{W}(\mathbf{D},x_{i})-\hat{u }(x_{i})\right|^{2}+w_{\rm reg}\sum_{i=1}^{N-1}\left|(\mathbf{D}_{i+1}-\mathbf {D}_{i})/h\right|^{2}$ $$

where $h$ is the spacing of the collocation points. Here we work with the vector $\mathbf{D}$ for simplicity. Althernatively, we can represent $D(x)$ as a neural network as in PINN and BiLO experiments.

We perform a grid search on the hyperparameters $l$ = 0.1, 0.2, 0.3, 0.4 and $w_{\rm reg}$=1e-3, 1e-4, 1e-5.
In Table 3, we show the 3 combinations of $l$ and $w_{\rm reg}$ with the best performance in terms of the $L_{2}$ error of the inferred $D(x)$ and the ground truth.
As shown in Table 3, the performance of the inference depends on properties of the pre-training dataset.
In practice, it might be difficult to know what does the ground truth unkown function look like.
This highlights the importance of the residual loss used in BiLO and PINN, which can help to learn the solution of the PDE without prior knowledge of the ground truth solution.

**Table 3: Comparison of BiLO and DeepONets ($l$ / $w_{\rm reg}$) pre-trained with datasets with different length scale $l$ and regularization weight $wreg$.**
| method | $\lVert D-D_{GT}\lVert_{\infty}$ | $\lVert D-D_{GT}\lVert_{2}$ | $\lVert u_{\rm NN}-u_{\rm FDM}\lVert_{\infty}$ | $\mathcal{L}_{\rm data}$ |
| --- | --- | --- | --- | --- |
| BiLO | 5.86e-2$\pm$1.99e-2 | 2.01e-2$\pm$7.94e-3 | 3.94e-3$\pm$1.93e-3 | 1.01e-4$\pm$1.80e-5 |
| DeepONet(0.2/1e-5) | 5.55e-2$\pm$7.99e-3 | 2.36e-2$\pm$2.05e-3 | 6.56e-3$\pm$2.36e-3 | 9.45e-5$\pm$1.47e-5 |
| DeepONet(0.4/1e-5) | 6.83e-2$\pm$2.76e-2 | 2.94e-2$\pm$1.14e-2 | 8.65e-3$\pm$9.37e-4 | 8.62e-5$\pm$1.36e-5 |
| DeepONet(0.4/1e-4) | 8.22e-2$\pm$2.08e-2 | 3.16e-2$\pm$8.63e-3 | 7.73e-3$\pm$1.89e-3 | 1.01e-4$\pm$1.83e-5 |

## Appendix E Additional Numerical Experiments

### E.1 Nonlinear System of ODEs

In this example, we consider inferring the unkown parameter $a$ in the following system of ODEs:

$$ $\begin{cases}\partialderivative{u_{1}}{t}=0.2u_{1}^{3}+4u_{2}^{3}\\ \partialderivative{u_{2}}{t}=au_{1}^{3}-0.2u_{2}^{3}\\ \end{cases}$ (33) $$

with initial condition $u_{1}(0)=1$, $u_{2}(0)=0$ and $t\in[0,1]$.
We set the ground truth $a_{\rm GT}=-4$ and take $a_{0}=-2$ as the initial guess.
To enforce the initial condition, the network have the form $u(t,a;W)=\mathcal{N}(t,a;W)t+u_{0}$, where $u_{0}$ is the initial condition (similar for the PINNs).
We set $|\mathcal{T}_{\rm res}|=|\mathcal{T}_{\rm dat}|=11$ and train for 20,000 steps.

##### Effect of residual-gradient loss

In Fig. 6, we visualize the function $u(t,a_{0}+\delta a;W)$ after pre-training with $a_{0}=-2$ for $\delta a$ = 0, -0.2, 0.2.
We can see that the neural network approximate the solution of the ODE for small variation of $a$.

Figure: Figure 6: Evaluating $u(t,a_{0}+\delta a;W)$ after pre-training with $a_{0}=-2$. The network approximates the solution of the ODE for small variation of $a$.
Refer to caption: extracted/5794873/fig_ode_var.png

##### Inference with Noise Data

In Table 4, we show the inference results for the ODE problem with noise $\epsilon\sim N(0,0.05)$.
We compare the performance of the PINNs with various $w_{\rm dat}$ and BiLO.
For the PINNs, as $w_{\rm dat}$ increase form 1e-3 to 1e1, the accuracy of the inferred parameter $a$ decreases and increases.
As $w_{\rm dat}$ goes above 1e-1, we can see the network tends to over-fit the data, as $\mathcal{L}_{\rm dat}$ gets smaller than the variance of the noise, and $\mathcal{L}_{\rm res}$ gets larger, while BiLO gives more accurate inferred parameter and ODE solution.
In Fig. 7, we show one instance of 5 random trials to compare the results of the PINN formulation with different pairs of weights $w_{\rm dat}$= 0.1, 1, 10, and BiLO.

**Table 4: Comparison of BiLO and PINNs (with various $w_{\rm dat}$) for a nonlinear ODE problem with noise ($\epsilon\sim N(0,0.04)$): mean (std) for predicted $a$ deviation from ground truth ($a_{\rm GT}=-4$) and losses $L_{r}$ and $L_{d}$**
| method | $\lvert a_{\rm pred}-a_{\rm GT}\lvert$ | $\mathcal{L}_{\rm data}$ | $\lVert u_{\rm NN}-u_{\rm FDM}\lVert_{\infty}$ |
| --- | --- | --- | --- |
| BiLO | 0.23$\pm$0.18 | 5.04e-2$\pm$1.28e-2 | 1.10e-1$\pm$8.31e-2 |
| PINN(1e-2) | 0.31$\pm$0.16 | 5.41e-2$\pm$1.05e-2 | 1.48e-1$\pm$7.60e-2 |
| PINN(1e-1) | 0.33$\pm$0.40 | 4.49e-2$\pm$1.13e-2 | 7.17e-2$\pm$5.81e-2 |
| PINN(1e0) | 1.39$\pm$0.76 | 1.71e-2$\pm$1.07e-2 | 3.37e-1$\pm$1.52e-1 |

Figure: Figure 7: Network predicted solutions $u_{\rm NN}$ (BiLO and PINNs with different $w_{\rm dat}$) and FDM solution $u_{\rm FDM}$ with the inferred parameter $a$.
Refer to caption: extracted/5794873/fig_ode_pred.png

### E.2 Infer the Initial Condition of a Heat Equation

In this example, we aim to infer the initial condition of a 1D heat equation from the final state.
Consider the heat equation

$$ $\begin{cases}u_{t}(x,t)=Du_{xx}(x,t)\\ u(x,0)=f(x)\\ u(0,t)=u(1,t)=0\\ \end{cases}$ (34) $$

on $x\in[0,1]$ and $t\in[0,1]$, with fixed diffusion coefficient $D=0.01$, and unknown initial condition $f(x)$, where $f(0)=f(1)=0$.
Our goal is to infer the initial condition $f(x)$ from observation of the final state $u(x,1)$.
We set the ground truth initial condition $f_{\rm GT}$ to be the hat function

$$ $f_{\rm GT}(x)=\begin{cases}2x,&\text{if }x\in[0,0.5)\\ 2-2x,&\text{if }x\in[0.5,1]\end{cases}$ (35) $$

We set the initial guess $f_{0}(x)=\sin(\pi x)$.
We can represent the unknown function $f(x;V)=s(\mathcal{N}(x;V))x(1-x)$, where $N_{f}$ is a fully connected neural network with 2 hidden layers and width 64, and $s$ is the softplus activation function (i.e., $s(x)=\log(1+\exp(x))$).
The transformation ensures that the initial condition satisfies the boundary condition and is non-negative.
For BiLO, the neural network is represented as $u(x,t,z)=N_{u}(x,t,z;W)x(1-x)t+z$, where $N_{u}$ is a fully connected neural network with 2 hidden layers and width 128.
For the PINN, we have $u(x,t;W,V)=N_{u}(x,t;W)x(1-x)t+f(x;V)$.
These transformations ensure that the networks satisfy the boundary and initial condition.

Let $X_{r}$, $X_{d}$ be spatial coordinates evenly spaced in $[0,1]$
and $T_{r}$ be temporal coordinates evenly spaced in $[0,1]$ (both including the boundary).
We set $\mathcal{T}_{\rm res}=X_{r}\times T_{r}$ and $|X_{r}|=|T_{r}|=51$. That is, the residual collocation points is a uniform grid in space and time.
We set $\mathcal{T}_{\rm dat}=X_{d}\times\{1\}$ and $|X_{d}|=11$. That is, the data collocation points is a uniform grid in space at the final time $t=1$.
We set the collocation point for the regularization loss of the unknown function $\mathcal{T}_{\rm reg}$ to be 101 evenly spaced points in the spatial domain.

To evaluate the performance of the inferred initial condition $f$,
we use the $L_{2}$ norm and the $L_{\infty}$ norm of the difference between the inferred initial condition and the ground truth initial condition,
which are evaluated at 1001 evenly spaced points in the spatial domain.

#### Without Noise

First we consider the case where the data is provided at $t=1$ without noise.
In this case, we also do not use regularization term for the initial condition.
In Fig. 8, and Table 5,
we show the results of PINNs various weights $w_{\rm dat}$= 0.1, 10, 1000, and BiLO.
We can see that BiLO achieved the best $e_{\rm 2}$ and $e_{\infty}$,
demonstrating the effectiveness in recovering the non-smooth initial condition.
With very large data loss, the error of the PINN increases.
This is because data is only provided at the final time, we need to solve the PDE accurately to infer the initial condition.

Figure: Figure 8: Predicted initial conditions of the heat equation (without noise) from 5 random seeds.
Refer to caption: extracted/5794873/fig_heat_nnz.png

**Table 5: Comparison of BiLO and PINNs (with various $w_{\rm dat}$) for inferring the unknown inititial condition (without noise), showing mean (std).**
| method | $\lVert f_{\rm NN}-f_{\rm GT}\lVert_{\infty}$ | $\lVert f_{\rm NN}-f_{\rm GT}\lVert_{2}$ | $\lVert u_{\rm NN}-u_{\rm FDM}\lVert_{\infty}$ | $\mathcal{L}_{\rm data}$ |
| --- | --- | --- | --- | --- |
| BiLO | 5.43e-2$\pm$1.00e-3 | 1.01e-4$\pm$7.46e-6 | 4.64e-4$\pm$2.57e-4 | 1.52e-9$\pm$5.89e-10 |
| PINN(1e-1) | 9.24e-2$\pm$2.21e-3 | 5.43e-4$\pm$2.81e-5 | 1.29e-3$\pm$1.53e-3 | 4.05e-6$\pm$4.57e-6 |
| PINN(1e1) | 8.69e-2$\pm$2.39e-3 | 4.31e-4$\pm$5.06e-5 | 2.44e-3$\pm$9.86e-4 | 1.62e-6$\pm$1.82e-6 |
| PINN(1e3) | 1.49e-1$\pm$4.19e-3 | 1.92e-3$\pm$1.34e-4 | 2.54e-2$\pm$2.92e-3 | 3.83e-8$\pm$5.36e-8 |

#### With Noise

In this experiment, we consider the case with noise $\epsilon\sim N(0,0.001)$.
Due to the ill-posedness of the inverse problem, we need to regularize the problem by the 2-norm of the derivative of the unknown function with $w_{\rm reg}=1e-2$.
In Fig. 9 and Table 6,
we show examples of the inferred initial condition and the PDE solution for the PINN formulation with various $w_{\rm dat}$.
In Table 6, for the PINN, we can see that as $w_{\rm dat}$ increase from 0.1 to 10, it seems that the reconstruction error decreases.
However, the $\mathcal{L}_{\rm dat}$ is becoming smaller than the variance of the noise, indicating that the PINN is overfitting the data.
This can also be observed from the Fig 9, for $w_{\rm dat}=1e3$ , we see larger discrepancy between $u_{\rm PINN}$ and $u_{\rm FDM}$.

Figure: Figure 9: Predicted initial condition $f(x)$ by BiLO and PINNs with various $w_{\rm dat}$.
Refer to caption: extracted/5794873/fig_heat_nz.png

**Table 6: Comparison of the BiLO and PINN (with various $w_{\rm dat}$) for a heat equation with unknown inititial condition (noise $\epsilon\sim N(0,0.001)$), showing mean (std).**
| method | $\lVert f_{\rm NN}-f_{\rm GT}\lVert_{\infty}$ | $\lVert f_{\rm NN}-f_{\rm GT}\lVert_{2}$ | $\lVert u_{\rm NN}-u_{\rm FDM}\lVert_{\infty}$ | $\mathcal{L}_{\rm data}$ |
| --- | --- | --- | --- | --- |
| BiLO | 2.41e-1$\pm$7.62e-3 | 6.11e-3$\pm$5.36e-4 | 1.46e-3$\pm$7.50e-4 | 4.08e-3$\pm$2.53e-4 |
| PINN(1e1) | 2.62e-1$\pm$2.22e-2 | 8.13e-3$\pm$2.79e-3 | 1.26e-1$\pm$3.27e-2 | 5.21e-4$\pm$1.55e-4 |
| PINN(1e2) | 2.53e-1$\pm$2.34e-2 | 7.11e-3$\pm$2.06e-3 | 1.38e-1$\pm$3.17e-2 | 2.61e-4$\pm$1.56e-4 |
| PINN(1e3) | 2.42e-1$\pm$4.65e-2 | 6.56e-3$\pm$2.70e-3 | 1.36e-1$\pm$3.33e-2 | 2.05e-4$\pm$1.65e-4 |

### E.3 2D Poisson Equation with Variable Diffusion Coefficient

The setup of this experiment is similar to the steady state Darcy flow inverse problem in [liPhysicsInformedNeuralOperator2024]. We consider the following 2D Poisson equation with variable diffusion coefficient in the unit square domain $\Omega=[0,1]\times[0,1]$ with Dirichlet boundary condition:

$$ $\begin{cases}-\nabla\cdot(A(\mathbf{x})\nabla u(\mathbf{x}))=f(\mathbf{x})& \text{in }\Omega\\ u(\mathbf{x})=0,\quad\text{on }\partial\Omega\end{cases}$ (36) $$

Our goal is to infer the variable diffusion coefficient $A(\mathbf{x})$ from the solution $u(\mathbf{x})$.

Let $\phi(\mathbf{x})$ be samples of a Gaussian random field (GRF) with mean 0 and squared exponential (Gaussian) covariance structure
$C(\mathbf{x},\mathbf{y})=\sigma\exp(-||\mathbf{x}-\mathbf{y}||^{2}/\lambda^{2})$, where the marginal standard deviation $\sigma=\sqrt{10}$ and the correlation length $l=0.01$ [constantineRandomFieldSimulation2024].
This GRF is different from [liPhysicsInformedNeuralOperator2024].
We generate the initial guess $A_{0}(\mathbf{x})=\text{sigmoid}(\phi_{0}(\mathbf{x}))\times 9+3$, where $\phi_{0}(\mathbf{x})$ is a sample of the GRF.
We consider the ground truth diffusion coefficient to be a piece-wise constant function:
$A_{\rm GT}(\mathbf{x})=12$ if $\phi_{\rm GT}(\mathbf{x})>0$ and $A_{\rm GT}(\mathbf{x})=3$ otherwise, where $\phi_{\rm GT}$ is another sample of the GRF.
The corresponding solution of $A_{\rm 0}$ and $A_{\rm GT}$ are denoted as $u_{0}$ and $u_{\rm GT}$.

We pretrain the BiLO with $A_{0}(\mathbf{x})$ and it’s corresponding solution $u_{0}(\mathbf{x})$ for 10,000 steps. And we fine-tune the BiLO for 5,000 steps using $u_{\rm GT}(\mathbf{x})$ to infer $A_{\rm GT}$. Following [liPhysicsInformedNeuralOperator2024], we use the total variation regularization $|\nabla A|$ with weight $w_{\rm reg}=1e-9$.

The unknown function is represendted $A(\mathbf{x};V)=s(\mathcal{N}(\mathbf{x};V))\times 9+3$, where $N_{f}$ is a fully connected neural network with 2 hidden layers and width 64, and $s$ is the sigmoid activation function (i.e., $s(u)=1/(1+\exp(-u))$).
The transformation is a smoothed approximation of the piece-wise constant function.
For BiLO, the neural network is represented as $u(\mathbf{x},z)=N_{u}(\mathbf{x},z;W)\mathbf{x}_{1}(1-\mathbf{x}_{1})\mathbf{x
}_{2}(1-\mathbf{x}_{2})$, where $N_{u}$ is a fully connected neural network with 2 hidden layers and width 128, and $z$ is our auxiliary variable such that $z=A(\mathbf{x};V)$.

Figure: Figure 10: Example 1 of inferring the variable diffusion coefficient. The relative l2 error of $u_{\rm NN}$ against $u_{\rm GT}$ is 1.3%. The thresholded (at the dashed line) inferred diffusion coefficient has classification accuracy of 98%
Refer to caption: extracted/5794873/fig_darcy_t4_s3.png

Figure: Figure 11: Example 2 of inferring the variable diffusion coefficient. The relative l2 error of $u_{\rm NN}$ against $u_{\rm GT}$ is 1.7%. The thresholded (at the dashed line) inferred diffusion coefficient has classification accuracy of 96%
Refer to caption: extracted/5794873/fig_darcy_t5_s4.png

In Figure 10 and Figure 11, we show two examples of the results, each with different initial guess $A_{0}$ and ground truth $A_{\rm GT}$.
In example 1 (see Figure 10), the relative error of the inferred diffusion coefficient is 1.3%. If we threshold the inferred diffusion coefficient at 7.5 (the mid-point of 3 and 12), the classification accuracy is 98%.
In example 2 (see Figure 11), the relative error of the inferred diffusion coefficient is 1.7%. If we threshold the inferred diffusion coefficient, the classification accuracy is 96%.
Our performance is comparable to the results (2.29% relative l2 error on $u$ and 97.10% classification accuracy) from the Physics-informed Neural Operator (PINO) in [liPhysicsInformedNeuralOperator2024], which require pretraining a FNO with synthetic dataset, and instance-wise fine-tuning with physics-informed loss.
For our method, we only need to pretrain the BiLO with a single initial guess. In addition, as shown in the figures, the intial guess can be very different from the ground truth.

## Appendix F Computational Cost

Compared with PINN, BiLO involve computing a higher order derivative term in the residual-gradient loss.
This increases the memory cost and computation time per step.
However, as shown in Fig. 4, BiLO might require fewer iterations to achieve certain accuracy of the parameters.

In Table. 7, we show the seconds-per-step and the maximum memory allocation of 1 run of BiLO and PINN for the various problems. The seconds per step is computed by total training time divided by the number of steps. The maximum memory allocation is the peak memory usage during the training. For for all the experiments, we use Quadro RTX 8000 GPU.
We note that the measured seconds-per-step is not subject to rigorous control as the GPU is shared with other users and many runs are performed simultaneously.
Detailed study of the computational efficiency of BiLO will be left for future work.

**Table 7: Example of computational cost of BiLO and PINN and their ratio for various problems.**
| Problem | Metric | BiLO | PINN | BiLO/PINN |
| --- | --- | --- | --- | --- |
| Fisher-KPP | sec-per-step | 0.13 | 0.086 | 1.51 |
| max-mem-alloc | 437 | 105 | 4.16 |  |
| Nonlinear ODE | sec-per-step | 0.05 | 0.024 | 2.08 |
| max-mem-alloc | 18.6 | 17.9 | 1.04 |  |
| 1D Poisson | sec-per-step | 0.05 | 0.03 | 1.67 |
| max-mem-alloc | 30 | 22 | 1.36 |  |
| Heat | sec-per-step | 0.096 | 0.06 | 1.60 |
| max-mem-alloc | 367 | 171 | 2.15 |  |

It is not straightforward to comparing the computational cost with Neural operators.
Neural operators can be very fast in the inference stage (solving inverse problem).
However, they have significant overhead, which involve preparing the training data, that is, solve the PDE numerically for a large collection of parameters, and pre-train the neural network.
The overall cost might be favorable in the many-query settings.
However, if we aim to solve the inverse problem once, the total computational cost might not be favorable.

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 6
presentation: 3
questions: 1. Discuss (perhaps using numerical evidence) that indeed the solutions computed are solutions to the PDE. Also discuss (again potentially with numerical evidence) that the solutions computed are local solutions and not global solutions.
2. Discuss (perhaps using numerical evidence) what the proposed algorithm actually converges to stationary solutions (or even global minima) of the bilevel problem.
soundness: 2
strengths: The problem considered is important and challenging. Across various communities do people search for efficient solutions, thus making the paper potentially significant. The paper is original in that the idea has not been considered before. The presentation is clear and the ideas do come across well. I really like the tumour growth example applying such methods to more than just box-standard applications.
summary: The paper investigates a method for parameter identification in PDEs (similar to optimal control). The idea is to replace the PDE constraint with novel functional that penalises both the residual as well as the derivative of the residual with respect to the parameters. The latter is local in the parameters, hence the motivation for the title of the paper. The paper also models the solution as a neural network similar to PINNS. This approach is equipped with an efficient learning strategy and initialisation. It is evaluated on various simple PDEs as well as a more complicated one related to tumour growth.
weaknesses: The paper is based on a lot of heuristics and I have doubts that this idea in general will be useful for others solving similar problems.

1. The entire idea of local solutions is neither well-defined nor do I believe that what is being computed are local solutions in the sense that they actually solve the PDE in a neighbourhood of the parameters. Also they still use the global residual encouraging both local and global behaviour. It is not clear how to entangle the two: are both important?

2. The optimisation algorithm in general will not solve the bilevel optimisation problem. This problem has been studied for decades (or more; see the papers cited in the paper itself) and the community agrees that the proposed algorithm does not do it (in general). I agree with the arguments in the appendix but this is for a much simpler case and in my opinion will not generalise. This is why methods like the "adjoint method" etc exist. Note that this is also a current research topic in the optimisation community which investigate similar but still very different algorithms compared to the one proposed here. The paper also makes the claim that the PDE constraint "does not need to be solved to solved to optimality". I agree this is true but the statement needs to quantitative. There will be an accuracy that is needed to make sure the computed directions actually descent. This is completely ignored at present. See e.g.
Pedregosa, Fabian. "Hyperparameter optimization with approximate gradient." International conference on machine learning. PMLR, 2016.
The paper starts with a very good initial estimation of the solution which is probably the key reason why it works well in the numerical examples. I don't believe this generalises well to other problems.

Minor comments:
- line 199, typo a(x, f(x))

---

confidence: 4
contribution: 2
initial_rating: 6
presentation: 3
questions: N.A.
soundness: 2
strengths: The paper basically argues for solving inverse problem with neural operators instead of PINN for both efficiency and accuracy. The support of the argument is mostly empirical. However, the experiment is well designed in both diversity and careful measurement, e.g., in addition to measuring accuracy of solution and unknow parameter (function), the author deliberately choose the initial guess that is far from ground truth in their experiment, which is a strong evidence for the effectiveness of the method.
summary: This paper intends to solve PDE constrained inverse problem with neural operators, which is formulated into a bi-level optimization problem. In the lower level optimization, a neural operator is learned and trained on a fixed initial guess of PDE parameter $\Theta_0$, which is named "local neural operator" by the authors. In the upper level optimization, the optimal or targeted parameter $\Theta$ is approximated.

In the experiment, the authors benchmarked their method on a tumor growth and imaging model and compared results to PINN-based method.
weaknesses: What is the motivation of using neural operator for PDE-constrained inverse problem? Neural operator is known for its generalization on various input functions, boundary conditions, initial data, etc. However, in the setting of this paper, a "local" neural operator is learned, meaning narrow generalization ability and abandoning the advantage of neural operators. Could the authors elaborate on the advantages of using a local neural operator versus a more general neural operator in this context? What specific benefits does the local approach provide for PDE inverse problems that outweigh the loss of broader generalization ability?

Meanwhile, PINN is known for low requirement of data and being suitable for inverse problem. The experiment indeed shows that the improvement of the method proposed here over PINN is quite marginal considering noise (see Table 2&3). Given the relatively small improvement over PINN shown in Tables 2 and 3, could the authors provide a more comprehensive discussion of the advantages of their method? Are there specific scenarios where the proposed approach significantly outperforms PINN that may not be captured in the current experiments?

Also, how does the proposed method compare to the bi-level optimization approach for PINNs presented in [1]? Could the authors discuss the key differences and potential advantages of their method over this existing work?

[1] Hao, Zhongkai, et al. "Bi-level Physics-Informed Neural Networks for PDE Constrained Optimization using Broyden's Hypergradients." The Eleventh International Conference on Learning Representations.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 2
questions: - On line 177, "The use of Lu0 is not mandatory for training the local operator with fixed Θ0, though it can speed up the training process", does it mean fixed initial parameters throughout training or just during initialization?
- Scalability. How does BiLO perform on higher-dimensional PDEs? How could BiLO reduce computational overhead for large-scale problems?
- Can the authors provide a more rigorous theoretical analysis of the convergence behavior for the bilevel optimization process?
- Comparison with SOTA Neural Operators. The paper focuses on comparing BiLO with PINNs. Could the authors provide some comparison over neural operators (e.g., DeepONet, FNO)?
soundness: 3
strengths: - **Elimination of Residual-Data-Loss Balancing**. By separating the optimization process into two levels, BiLO eliminates the need to balance the residual and data loss;
- **Local Operator Learning**. Local PDE solution operators ensure precise gradient computation and lead to faster parameter inference.
- **Robustness**. Experimental results show that BiLO can handle noisy measurements and unseen data. 
- **Generalization Power**. The proposed method could infer unknown functions (e.g., variable diffusion coefficients) with the auxiliary variables.
summary: This paper proposes BiLO, which formulates upper and lower inverse problems into bilevel optimization. BiLO integrates adjoint methods and PINNs, while eliminating the need for balancing residual and data loss. BiLO has high robustness to sparse and noisy data, and effectively infers unknown PDE parameters and functions through auxiliary variables. Experimental results show that BiLO enforces strong PDE constraints, is robust to sparse and noisy data, and eliminates the need to balance the residual and the data loss, which is inherent to the soft PDE constraints in many existing methods.
weaknesses: - **Lack of Theoretical Guarantees**. The paper offers mainly empirical results. 
- **Scalability Concerns**. The experiments are limited to low-dimensional PDE problems. It is unclear how the method scales to higher-dimensional or more complex PDEs.
- **Computational Overhead**. Although the method achieves accurate results, solving both upper and lower-level optimization problems simultaneously introduces computational complexity.
- **Presentation**. The presentation of this paper could be further improved. The paper also seems to be completed in a rush and needs further proofreading. There are multiple typos (e.g., line 37, "... or deep learning, methods.", redundant comma) and duplicate bibitems (e.g., second and third last references on pp.15 are identical).

---

confidence: 5
contribution: 3
initial_rating: 6
presentation: 4
questions: 1. The initial parameter $\Theta$ seems important to the propsoed method. How do you choose it?
2. Again, I wonder how the proposed method works if the initial guess of $\Theta$ is far away from the ground truth.
3. In this work, the loss function is easy to construct since the equality constraint only exists. How do you generalize your method to inequality constraints? The inequality constraints are common in practice.
4. Which grids does the observed data locate? Please clarify this issue in the numerical experiments.
5. In this work, the noise is set to the Gaussian with a small variance. How about the noise with a large magnitude? 
6. About the training procedure. The simultaneous gradient descent is applied in the fine-tuning step instead of the alternative gradient descent. Could you plot the training loss curve for this step?
soundness: 3
strengths: 1. This paper is well-written. 
2. The proposed method is mesh-free, avoiding the difficulties caused by mesh generation. 
3. The algorithm is easy to implement and can provide fast inference for inverse problems once it has finished training.
summary: This work aims to infer parameters in parametric PDEs using some observed solution data. This type of problem is widely found in scientific and engineering applications. The authors propose a neural network method for solving inverse problems governed by partial differential equations (PDEs). This method contains two steps: the pre-train step using the data and collocation points and the fine-tune step using the automatic differentiation. The proposed approach incorporates the traits of PINN, operator learning, and the adjoint method.
weaknesses: 1. The PDE-constrained optimization problem considered in this work only involves the equality constraint, but in practice, the inequality constraints are typical, e.g., the box constraint. 
2. The pre-train step is crucial, and this step is minimizing the residual loss with collocation points. When the solution has a low-regularity property, this pre-train step will have a large generalization error if the collocation points are not properly chosen. This will overshadow the capability of the proposed method.

Meta-review: This paper introduces a new approach for solving inverse problems for partial differential equations (PDEs). Specifically, the approach formulates the problem as a bilevel optimization problem. This efficient approach is also able to  enforce PDE constraints. Results demonstrate good performance on a range of different problems, including sparse and noisy data. 

The reviewers agree that this work is novel and original. Moreover, the experiments are interesting. Also, the the paper is well written. However, there are also some concerns. Reviewers are skeptical how useful this framework will be in general. The experimental studies are limited to DeepONet. It would be good to see more experiments for different model architectures. There is lack of supporting theory for this approach, which is uncommon for optimization oriented ML papers.

While this is an overall nice paper, all reviewers indicated concerns in the final discussion. There is no strong advocate for this paper. I feel that the authors need to demonstrate that the proposed method is applicable more generally, either by providing a more extensive set of experiments that includes different architectures and a more diverse set of problems, or through more rigorous theory that supports the optimization approach. Thus, I recommend to reject this paper in its current form.

Decision: Reject