Bayesian BiLO: Bilevel Local Operator Learning for Efficient Uncertainty Quantification of Bayesian PDE Inverse Problems with Low-Rank Adaptation

TL;DR

This work addresses uncertainty quantification in PDE inverse problems by marrying gradient-based Bayesian inference with a bilevel local operator learning framework. The upper level performs posterior sampling of PDE parameters $\theta$ via Hamiltonian Monte Carlo, while the lower level deterministically learns a local solution operator $u(x,\theta;W)$ using a LoRA-enhanced neural network to enforce PDE constraints through the local operator loss $\mathcal{L}_{LO}$. The authors prove that inexact lower-level solves induce only $O(\epsilon)$-level errors in gradients and posterior accuracy, and demonstrate the method's efficiency and accuracy across nonlinear Poisson, GBM tumor growth, stochastic-rate inference, and Darcy flow problems, with LoRA providing substantial speedups especially for larger networks. Compared to BPINNs, B-BiLO avoids sampling high-dimensional neural weights and mitigates ill-conditioning from PDE residuals, offering a scalable and robust approach for uncertainty quantification in complex PDE models with potential for extension to 3D problems and high-dimensional parameter spaces.

Abstract

Uncertainty quantification in PDE inverse problems is essential in many applications. Scientific machine learning and AI enable data-driven learning of model components while preserving physical structure, and provide the scalability and adaptability needed for emerging imaging technologies and clinical insights. We develop a Bilevel Local Operator Learning framework for Bayesian inference in PDEs (B-BiLO). At the upper level, we sample parameters from the posterior via Hamiltonian Monte Carlo, while at the lower level we fine-tune a neural network via low-rank adaptation (LoRA) to approximate the solution operator locally. B-BiLO enables efficient gradient-based sampling without synthetic data or adjoint equations and avoids sampling in high-dimensional weight space, as in Bayesian neural networks, by optimizing weights deterministically. We analyze errors from approximate lower-level optimization and establish their impact on posterior accuracy. Numerical experiments across PDE models, including tumor growth, demonstrate that B-BiLO achieves accurate and efficient uncertainty quantification.

Figures (6)

• Figure 1: A schematic for solving the Bayesian PDE inverse problem using BiLO and LoRA. (A) At the upper level, we sample the PDE parameters $\theta$ from the posterior distribution using gradient-based MCMC. (B) The full operator $u(x, \theta)$ (gray surface) and its local approximation $u(x, \theta_1)$ and $u(x, \theta_2)$ (blue surface), which satisfies vanishing residual and residual-gradient conditions. The local operator suffices to compute the gradient of the potential energy $U(\theta)$ at $\theta_1$ and $\theta_2$. (C) At the lower level, we train the neural network to approximate the local operator (blue surface) at different $\theta$. As $\theta$ changes, only low rank adaptation (LoRA) is used to update the neural network weights $W$.
• Figure 2: Inference results for the nonlinear Poisson problem. (a) Noisy data, GT solution, mean and std of inferred $u$ using B-BiLO (b) results using BPINN with $\sigma_f=0.1$. (c) Posterior distribution of the PDE parameter $k$ using BiLO and BPINN with different $\sigma_f$, compared with the reference (Metropolis-Hastings with numerical solver). (d) Accuracy-efficiency comparison: x-axis shows the -log of the std. at the orgin, larger values indicate better physical uncertainty quantification. y-axis shows the effective sample size (ESS) per second of wall time, larger values indicate more efficient sampling. As $\sigma_f$ decreases, BPINN becomes more accurate but less efficient. Overall B-BiLO is more accurate and efficient at uncertainty quantification.
• Figure 3: Accuracy and efficiency of B-BiLO with LoRA of different ranks. (a) Comparison of the posterior distribution of $k$ obtained by B-BiLO with LoRA of different ranks and Full FT. Diamond markers denote the MAP estimates, and bars indicate the 95% highest-density posterior region (HDPR). (b) Relative efficiency of LoRA (rank 8) as ratios with respect to Full FT. The x-axis shows the number of neurons per layer; the left y-axis (blue line) shows the relative average time per sample, and the right y-axis shows the relative maximum memory usage during sampling.
• Figure 4: Inference of patient-specific GBM growth parameters using B-BiLO with LoRA rank 8. (a) Posterior distribution of the parameters $D, \rho, u_c^{\rm WT}, u_c^{\rm TC}$ and the prior (dashed), compared with Full FT. (b) Posterior mean of the tumor cell density $u$. (c) Ground truth MRI data with segmented WT and TC regions (filled), the MAP predicted segmentation (contour), and the predicted infiltration margin (where $u=1\%$). (d) Posterior standard deviation of the tumor cell density $u$.
• Figure 5: Inference of gene expression parameters $\lambda$ (birth rate) and $\mu$ (degradation rate) using B-BiLO. (a): The 90% highest-density posterior region (HDPR) and maximum a posteriori (MAP) estimates of the solution obtained using B-BiLO with LoRA rank 4, Full FT, and the reference method (MH with numerical solver). Inset: 90% HDPR and MAP of the solution $u(x)$. (b): Ground truth (GT), prior, and posterior distributions for $\lambda$ and $\mu$, with posterior sampled via B-BiLO using LoRA rank 4.