Bilevel Local Operator Learning (BiLO)

• BiLO is a neural network framework that formulates PDE inverse problems as a bilevel optimization, combining data loss minimization and local operator training.
• It employs a locally trained surrogate model to accurately approximate the PDE solution operator, facilitating robust and efficient descent direction estimation.
• The Bayesian extension integrates LoRA for low-rank adaptation, enabling fast posterior sampling and calibrated uncertainty quantification in complex inverse problems.

Bilevel Local Operator Learning (BiLO) is a 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, it minimizes the data loss with respect to the PDE parameters. At the lower level, it trains 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. In the second part of the series, the same framework is extended to the Bayesian inference framework: at the lower level, a network approximates the local solution operator by minimizing the local operator loss with respect to the weights of the neural network, and at the upper level, the PDE parameters are sampled from the posterior distribution (Zhang et al., 2024, Zhang et al., 22 Jul 2025).

1. Deterministic bilevel formulation

In the deterministic formulation, the spatial, or spatial-temporal, domain is denoted by $\Omega \subset \mathbb{R}^d$, the data-collocation points by $T$, and the observed data by $u_{\mathrm{obs}}(x)$ at $x \in T$. The PDE is assumed to take the form

$F(D^k u(x), \ldots, Du(x), u(x), \Theta)=0,$

where $\Theta \in \mathbb{R}^p$ are the unknown PDE parameters. BiLO introduces a neural network $u(x,\Theta;W)$ parameterized by weights $W$ to locally approximate the solution operator $u(x;\Theta)$ in a neighborhood of the current $\Theta$ (Zhang et al., 2024).

The upper-level loss is the data-mismatch

$T$0

The lower-level construction begins with the pointwise residual

$T$1

It then penalizes both the squared residual and its sensitivity to $T$2:

$T$3

$T$4

The local operator loss is

$T$5

where $T$6 is a small constant, for example $T$7, that controls the strength of the residual-gradient term.

The bilevel problem is

$T$8

The stated role of the lower level is not to produce a global neural surrogate over the full parameter space, but to provide a local surrogate accurate enough to supply a descent direction for the upper-level optimization.

2. Joint gradient descent and the local-operator mechanism

BiLO solves the two levels simultaneously by gradient descent rather than an inner-outer loop because $T$9 need only be accurate enough to give a descent direction $u_{\mathrm{obs}}(x)$0. The scalar-parameter algorithm takes as input data collocation sets $u_{\mathrm{obs}}(x)$1 and $u_{\mathrm{obs}}(x)$2, an initial PDE guess $u_{\mathrm{obs}}(x)$3, network weights $u_{\mathrm{obs}}(x)$4, a pretraining solution $u_{\mathrm{obs}}(x)$5 at $u_{\mathrm{obs}}(x)$6 as an optional input, and learning rates $u_{\mathrm{obs}}(x)$7 and $u_{\mathrm{obs}}(x)$8 (Zhang et al., 2024).

Optional pretraining minimizes

$u_{\mathrm{obs}}(x)$9

with

$x \in T$0

Fine-tuning then applies simultaneous steps for $x \in T$1:

1. compute $x \in T$2 via autodiff through $x \in T$3;
2. compute $x \in T$4;
3. update

$x \in T$5

$x \in T$6

In practice, the reported implementation uses ADAM with $x \in T$7 and $x \in T$8. The implementation details further specify pretraining $x \in T$9 for $F(D^k u(x), \ldots, Du(x), u(x), \Theta)=0,$0k steps on $F(D^k u(x), \ldots, Du(x), u(x), \Theta)=0,$1 and fine-tuning for $F(D^k u(x), \ldots, Du(x), u(x), \Theta)=0,$2k steps by joint GD. The explicit methodological claim is that BiLO enforces strong PDE constraints and eliminates the need to balance the residual and the data loss, which is inherent to the soft PDE constraints in many existing methods.

3. Network architecture and enforcement of boundary or initial conditions

For scalar $F(D^k u(x), \ldots, Du(x), u(x), \Theta)=0,$3, the reported architecture begins with an embedding layer

$F(D^k u(x), \ldots, Du(x), u(x), \Theta)=0,$4

where $F(D^k u(x), \ldots, Du(x), u(x), \Theta)=0,$5, $F(D^k u(x), \ldots, Du(x), u(x), \Theta)=0,$6 is a fixed random matrix, $F(D^k u(x), \ldots, Du(x), u(x), \Theta)=0,$7, and $F(D^k u(x), \ldots, Du(x), u(x), \Theta)=0,$8. The hidden portion consists of $F(D^k u(x), \ldots, Du(x), u(x), \Theta)=0,$9 fully connected layers of width $\Theta \in \mathbb{R}^p$0 with $\Theta \in \mathbb{R}^p$1 activations, and the output layer is a single scalar $\Theta \in \mathbb{R}^p$2 (Zhang et al., 2024).

Boundary-condition enforcement is handled by output transformations. For Dirichlet data $\Theta \in \mathbb{R}^p$3 on $\Theta \in \mathbb{R}^p$4, the construction

$\Theta \in \mathbb{R}^p$5

is used, where $\Theta \in \mathbb{R}^p$6 vanishes on $\Theta \in \mathbb{R}^p$7 and $\Theta \in \mathbb{R}^p$8 enforces the boundary or initial data. The implementation notes specify enforcing BC/IC via output transformations as in Section 3.1 of the paper. For unknown functions constrained by positivity, the activation is $\Theta \in \mathbb{R}^p$9 for the $u(x,\Theta;W)$0-network and softplus when enforcing $u(x,\Theta;W)$1.

The reported collocation choices illustrate the discretization used in the experiments: Fisher–KPP uses $u(x,\Theta;W)$2 on a grid in $u(x,\Theta;W)$3 and $u(x,\Theta;W)$4 points at $u(x,\Theta;W)$5; the Poisson problem uses $u(x,\Theta;W)$6 points in $u(x,\Theta;W)$7; and the heat equation uses $u(x,\Theta;W)$8 in $u(x,\Theta;W)$9 and $W$0 at $W$1.

4. Inference of unknown functions through an auxiliary variable

When the PDE depends on an unknown function $W$2, BiLO parametrizes $W$3 by another network $W$4, but avoids treating $W$5 as an input to $W$6 directly because it would be too high-dimensional. Instead, it introduces an auxiliary field

$W$7

The residual is then defined with this auxiliary variable:

$W$8

The data loss becomes

$W$9

and regularization on $u(x;\Theta)$0 is

$u(x;\Theta)$1

or any Tikhonov penalty (Zhang et al., 2024).

The lower-level loss is

$u(x;\Theta)$2

with

$u(x;\Theta)$3

$u(x;\Theta)$4

The resulting bilevel problem is

$u(x;\Theta)$5

and is again solved by joint gradient descent on $u(x;\Theta)$6.

The heat-equation example is the principal illustration of this mechanism: BiLO is reported to infer an unknown initial condition from the final state and to recover non-smooth $u(x;\Theta)$7 robustly, while PINN requires extremely large or small $u(x;\Theta)$8 to even converge.

5. Reported behavior on inverse problems and relation to PINNs

The central comparison in the deterministic paper is between BiLO and PINNs with soft-constraint objective $u(x;\Theta)$9. The stated advantage is strong PDE fidelity: unlike PINNs’ soft-constraint $\Theta$0, BiLO enforces zero residual and zero residual-gradient sufficiently tightly in the lower level, eliminating delicate $\Theta$1-tuning. The other principal claim is robustness to sparse and noisy data (Zhang et al., 2024).

The Fisher–KPP reaction–diffusion experiment infers $\Theta$2 from single-time data with $\Theta$3–$\Theta$4 noise; the reported result is that BiLO recovers $\Theta$5 with $\Theta$6 error, while PINN degrades heavily unless $\Theta$7 is hand-tuned. In the $\Theta$8D Poisson problem with variable diffusion $\Theta$9, BiLO is reported to capture the “hat” shape and kink, whereas PINN either oversmooths or oscillates depending on $T$00. In the nonlinear ODE system that infers $T$01 in $T$02, BiLO is described as stable under noise $T$03, while PINN fails for most $T$04.

These experiments define the framework’s practical positioning. The paper also lists suggested extensions for new PDE inverse problems: increasing network width or depth for higher-dimensional domains, combining adaptive collocation or importance sampling for $T$05, leveraging multi-fidelity pretraining if a rough solver is available, and extending to time-dependent unknowns $T$06 by including $T$07 in the $T$08-embedding. These are presented as suggestions rather than established experimental conclusions.

6. Bayesian extension, posterior sampling, and low-rank adaptation

Part II extends BiLO to the Bayesian inference framework for PDE-constrained optimization problems. The finite-dimensional parameter is written as $T$09, the forward PDE is

$T$10

and the local solution operator is approximated by

$T$11

where $T$12 are the network weights. The bilevel formulation is

$T$13

with lower-level loss

$T$14

and posterior

$T$15

where a prior $T$16 and data-misfit

$T$17

are assumed (Zhang et al., 22 Jul 2025).

The upper-level Bayesian inference is solved using a gradient-based MCMC, for example Preconditioned Crank–Nicolson (pCN) or Hamiltonian Monte Carlo (HMC). At each proposed parameter $T$18, the lower-level problem in $T$19 is solved approximately to tolerance $T$20, after which $T$21 and $T$22 are computed. The gradient uses $T$23, obtained by automatic differentiation through the network and the PDE residual term. The combined algorithm is given as BiLO–LoRA–HMC.

To speed up repeated lower-level solves for varying $T$24, Part II introduces LoRA (Low-Rank Adaptation). Let $T$25 be the pre-trained network at some reference $T$26. Updates are restricted to

$T$27

At each MCMC step, only $T$28 are optimized to solve the lower-level problem to tolerance $T$29. The number of trainable parameters per leapfrog step is therefore reduced from $T$30 to $T$31 with $T$32, and the paper states that this drastically cuts both memory and computation.

7. Error analysis and reported uncertainty-quantification results

Part II analyzes the effect of inexact lower-level optimization on both the sampler dynamics and the posterior distribution. Let $T$33 denote the exact lower-level minimizer with $T$34, and let $T$35 be an inexact solution with $T$36. Under mild assumptions, specifically PDE stability and Lipschitz bounds, the paper states two principal results: a dynamic error in the hypergradient,

$T$37

and a static error in the posterior,

$T$38

The stated implication is that solving the lower-level problem to tolerance $T$39 perturbs both the MCMC dynamics and the target posterior only by $T$40 (Zhang et al., 22 Jul 2025).

The numerical experiments are organized around three PDE inverse problems. For nonlinear Poisson, $T$41, BiLO recovers $T$42 and $T$43 with accuracy matching a fine-grid MH reference; LoRA reduces per-step cost by factors $T$44 with negligible loss in ESS or accuracy; and pCN and HMC acceptance rates are $T$45 with stepsizes unreachable by BPINNs. For the $T$46D point-process PDE, $T$47, inference of $T$48 from point data yields $T$49 MAP error, and BiLO-LoRA with $T$50 achieves MOSI comparable to full fine-tuning at $T$51. For $T$52D Darcy flow, $T$53, the framework infers $T$54 KL-expansion parameters of $T$55 from noisy observations of $T$56; BiLO + LoRA with $T$57 converges $T$58 faster than full fine-tuning for width $T$59k networks, with $T$60 $T$61 error and accurate $T$62 confidence intervals.

Across these examples, the paper states that BiLO–LoRA–HMC attains high-fidelity posterior inference, large-stepsize stability unlike standard BPINNs, speedups of $T$63–$T$64, and calibrated UQ. Taken together with Part I, this places BiLO as a bilevel framework whose defining feature is local operator learning at the lower level and parameter inference, deterministic or Bayesian, at the upper level.