stat.MLMar 16, 2026

Preconditioned One-Step Generative Modeling for Bayesian Inverse Problems in Function Spaces

Authors: Zilan Cheng, Li-Lian Wang, Zhongjian Wang

Organizations: Division of Mathematical Sciences, School of Physical and Mathematical Sciences, Nanyang Technological University, 637371, Singapore.

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

Abstract

We propose a machine-learning algorithm for Bayesian inverse problems in the function-space regime. Based on one-step generative transport, the method learns an amortized neural operator whose pushforward of a Gaussian source approximates the posterior distribution conditioned on each new observation. We show that white-noise sources are incompatible with the function-space limit, and therefore adopt a prior-aligned GRF as the source. We justify this choice through the Lipschitz regularity of the resulting one-step conditional posterior transport and numerical experiments on linear inverse and PDE-based inverse problems. The method is not distilled from MCMC: it is trained only with prior samples and simulated partial noisy observations. Once trained, it generates a 64×6464\times64 posterior sample in ∼10−3\sim 10^{-3}s, avoiding repeated forward-model evaluations in MCMC and repeated network evaluations in multistep generative samplers while matching key posterior summaries.

Figures & tables

Figure 1: Conditional neural-operator parameterization of the mean-flow predictor. A source draw ξ∼ρ\xi\sim\rho is lifted by P\mathcal{P} , passed through LL conditional operator layers, and projected by Q\mathcal{Q} to produce the mean-flow correction. During training, time pairs 0≤r≤t≤10\leq r\leq t\leq 1 are sampled randomly; at inference, the learned predictor is evaluated at (r,t)=(0,1)(r,t)=(0,1) to generate one-step posterior samples.
Figure 2: Linear inverse problem. The first row shows the input field, sparse observations, reference posterior mean, predicted posterior mean, and mean error. The second row shows representative posterior samples generated by the proposed sampler.
Figure 3: Energy spectra E(k)E(k) of the posterior distributions for the 65×6565\times 65 linear inverse problem under different choices of source kernels.
Sourceϵμeval\epsilon_{\mu}^{\mathrm{eval}}ϵσeval\epsilon_{\sigma}^{\mathrm{eval}}
ρC\rho_{C}4.94×10−24.94\times 10^{-2}2.72×10−22.72\times 10^{-2}
ρW\rho_{W}8.24×10−28.24\times 10^{-2}4.13×10−24.13\times 10^{-2}
ρ1\rho_{1}6.28×10−26.28\times 10^{-2}2.49×10−22.49\times 10^{-2}
ρ2\rho_{2}8.15×10−28.15\times 10^{-2}8.33×10−28.33\times 10^{-2}
Table 1: Effect of the source measure on posterior approximation for the linear inverse problem.
Figure 4: Darcy inverse problem. The first row shows the unknown field, sparse observations, reference posterior mean, predicted posterior mean, and mean error. The second row shows representative posterior samples generated by the proposed sampler.
Figure 5: Advection inverse problem. The first row shows the unknown field, sparse observations, reference posterior mean, predicted posterior mean, and mean error. The second row shows representative posterior samples generated by the proposed sampler.
Figure 6: Reaction-diffusion inverse problem. The first row shows the unknown field, sparse observations, reference posterior mean, predicted posterior mean, and mean error. The second row shows representative posterior samples generated by the proposed sampler.
Figure 7: Navier-Stokes inverse problem. The first row shows the unknown field, sparse observations, reference posterior mean, predicted posterior mean, and mean error. The second row shows representative posterior samples generated by the proposed sampler.
SamplerStepsSampling time (s/sample)
MCMC (Darcy)3×1043\times 10^{4}2.05×1012.05\times 10^{1}
MCMC (Advection)3×1043\times 10^{4}7.01
MCMC (Reaction–Diffusion)3×1043\times 10^{4}1.04×1011.04\times 10^{1}
MCMC (Navier–Stokes)3×1043\times 10^{4}1.97×1011.97\times 10^{1}
Multistep diffusion3203201.05×10−11.05\times 10^{-1}
Multistep flow matching3203201.67×10−11.67\times 10^{-1}
Table 2: Wall-clock generation time for producing posterior samples conditioned on a single fixed observation yobsy_{\mathrm{obs}} . Hardware: MCMC is run in parallel on 1616 cores of the CPU ( 3232 thread at 3.63.6 GHz) cores; learned samplers are evaluated on a single NVIDIA RTX A6000 GPU. The per-sample clock time is computed by averaging over the generation time of 10001000 posterior samples.
DatasetSourceϵμeval\epsilon_{\mu}^{\mathrm{eval}}ϵσeval\epsilon_{\sigma}^{\mathrm{eval}}
DarcyAnisotropic5.55×10−25.55\times 10^{-2}4.27×10−24.27\times 10^{-2}
White8.26×10−28.26\times 10^{-2}4.30×10−24.30\times 10^{-2}
AdvectionAnisotropic6.08×10−26.08\times 10^{-2}5.09×10−25.09\times 10^{-2}
White1.63×10−11.63\times 10^{-1}5.14×10−25.14\times 10^{-2}
Reaction– DiffusionAnisotropic7.62×10−27.62\times 10^{-2}8.68×10−28.68\times 10^{-2}
White7.89×10−27.89\times 10^{-2}8.91×10−28.91\times 10^{-2}
Table 3: Posterior approximation errors of the conditional one-step sampler on PDE inverse problems.
Appendix figures & tables9 assets

Supplementary material from the paper’s appendix.

Appendix

DatasetSourceϵμeval\epsilon_{\mu}^{\mathrm{eval}}ϵσeval\epsilon_{\sigma}^{\mathrm{eval}}
DarcyAnisotropic5.55×10−25.55\times 10^{-2}5.45×10−25.45\times 10^{-2}
White7.56×10−27.56\times 10^{-2}4.37×10−24.37\times 10^{-2}
AdvectionAnisotropic6.36×10−26.36\times 10^{-2}3.76×10−23.76\times 10^{-2}
White1.43×10−11.43\times 10^{-1}5.08×10−25.08\times 10^{-2}
Reaction–DiffusionAnisotropic7.23×10−27.23\times 10^{-2}7.94×10−27.94\times 10^{-2}
White8.08×10−28.08\times 10^{-2}7.85×10−27.85\times 10^{-2}
Appendix
Table 4: Posterior approximation errors of the PODNO-parameterized Mean-Flow sampler.
DatasetSourceϵμeval\epsilon_{\mu}^{\mathrm{eval}}ϵσeval\epsilon_{\sigma}^{\mathrm{eval}}
DarcyAnisotropic1.02×10−11.02\times 10^{-1}1.48×10−11.48\times 10^{-1}
White1.41×10−11.41\times 10^{-1}1.02×10−11.02\times 10^{-1}
AdvectionAnisotropic7.34×10−27.34\times 10^{-2}1.37×10−11.37\times 10^{-1}
White2.04×10−12.04\times 10^{-1}1.04×10−11.04\times 10^{-1}
Reaction–DiffusionAnisotropic9.17×10−29.17\times 10^{-2}2.04×10−12.04\times 10^{-1}
White1.44×10−11.44\times 10^{-1}1.44×10−11.44\times 10^{-1}
Appendix
Table 5: Posterior approximation errors of the U-Net-parameterized Mean-Flow sampler.
DatasetSourceϵμeval\epsilon_{\mu}^{\mathrm{eval}}ϵσeval\epsilon_{\sigma}^{\mathrm{eval}}
DarcyAnisotropic5.43×10−25.43\times 10^{-2}6.16×10−26.16\times 10^{-2}
White2.80×1092.80\times 10^{9}3.14×1093.14\times 10^{9}
AdvectionAnisotropic7.19×10−27.19\times 10^{-2}5.05×10−25.05\times 10^{-2}
White1.11×10101.11\times 10^{10}1.91×1051.91\times 10^{5}
Reaction-DiffusionAnisotropic8.23×10−28.23\times 10^{-2}9.64×10−29.64\times 10^{-2}
White2.16×10102.16\times 10^{10}6.14×1056.14\times 10^{5}
Appendix
Table 6: Posterior approximation errors of the improved Mean-Flow sampler.
DatasetSourceϵμeval\epsilon_{\mu}^{\mathrm{eval}}ϵσeval\epsilon_{\sigma}^{\mathrm{eval}}
DarcyAnisotropic1.08×10−11.08\times 10^{-1}7.00×10−27.00\times 10^{-2}
White1.19×10−11.19\times 10^{-1}7.71×10−27.71\times 10^{-2}
AdvectionAnisotropic9.90×10−29.90\times 10^{-2}4.30×10−24.30\times 10^{-2}
White2.99×10−12.99\times 10^{-1}8.24×10−28.24\times 10^{-2}
Reaction-DiffusionAnisotropic1.08×10−11.08\times 10^{-1}1.02×10−11.02\times 10^{-1}
White1.16×10−11.16\times 10^{-1}9.21×10−29.21\times 10^{-2}
Appendix
Table 7: Posterior approximation errors of the multistep diffusion baseline.
Figure 8: Effect of the number of test-time sampling steps on posterior accuracy for the linear inverse problem for the multistep diffusion model. We use the same model trained with 320 sampling steps and vary only the number of sampling steps at inference time.
Figure 9: Comparison of posterior statistics from our model and DPS inference on the same 65×6565\times 65 linear test instance as in Figure 2 . DPS uses an unconditionally trained diffusion model with a guidance weight, and we report the best-performing result obtained by tuning this weight over [0,1][0,1] , attained at 0.10.1 .
Figure 10: Comparison of posterior statistics from our model and DPS inference on the same 65×6565\times 65 darcy test instance as in Figure 4 . DPS uses an unconditionally trained diffusion model with a guidance weight, and we report the best-performing result obtained by tuning this weight over [0,1][0,1] , attained at 0.10.1 .
SymbolDescription
Ω\OmegaSpatial domain
ddSpatial dimension
ssSobolev regularity exponent defining the distribution space H−s\mathcal{H}^{-s}
X,x,X~X,x,\widetilde{X}Unknown random function, its realization, and encoded variable
Yobs,yobs,Y~obsY_{\mathrm{obs}},y_{\mathrm{obs}},\widetilde{Y}_{\mathrm{obs}}Observation random variable, its realization, and encoded observation
Ξ,ξ\Xi,\xiSource random variable and its realization, sampled from ρ\rho
Appendix
Table 8: Notation.
SettingSymbolValue
FNO Fourier modesKK12
FNO layersLL4
Network widthNwidthN_{\mathrm{width}}96
Learning ratelr\mathrm{lr}10−410^{-4}
Flow ratioνflow\nu_{\mathrm{flow}}0.25
Batch sizeNbatchN_{\mathrm{batch}}20
Appendix
Table 9: Network architecture and training hyperparameters.
More in this topic

Bayesian Inverse Problems

Bayesian inverse problems infer hidden signals or physical parameters from indirect, noisy measurements while quantifying uncertainty.

See all 118 papers

Explore similar work

CardsList

  1. Amortized Energy-Based Bayesian Inference

    May 14, 2026Hojjat Kaveh, Ricardo Baptista, Andrew M. StuartAmortized InferenceBayesian Inverse Problems

    California Institute of Technology · University of Toronto

  2. A Gibbs posterior sampler for inverse problem based on prior diffusion model

    Feb 11, 2026Jean-François GiovannelliMarkov Chain Monte CarloBayesian Inverse Problems

    Groupe Signal-Image, IMS (Univ. Bordeaux, CNRS, BINP), Talence, France

  3. Flow Annealing Posterior Sampling for Function-Space Regression and Inverse Problems

    Jun 21, 2026Yaozhong Shi, Zachary E. Ross, Yisong YueBayesian Inverse ProblemsPosterior Sampling

    California Institute of Technology