Conditional neural control variates for variance reduction in Bayesian inverse problems

Ali Siahkoohi; Hyunwoo Oh

stat.MLFeb 24, 2026

Conditional neural control variates for variance reduction in Bayesian inverse problems

Authors: Ali Siahkoohi, Hyunwoo Oh

Organizations: Department of Computer Science, University of Central Florida, Orlando, Florida, USA · Department of Physics and Maryland Center for Fundamental Physics, University of Maryland, College Park, Maryland, USA

Abstract

Bayesian inference for inverse problems involves computing expectations under posterior distributions--e.g., posterior means, variances, or predictive quantities--typically via Monte Carlo (MC) estimation. When the quantity of interest varies significantly under the posterior, accurate estimates demand many samples--a cost often prohibitive for partial differential equation-constrained problems. To address this challenge, we introduce conditional neural control variates, a modular method that learns amortized control variates from joint model-data samples to reduce the variance of MC estimators. To scale to high-dimensional problems, we leverage Stein's identity to design an architecture based on an ensemble of hierarchical coupling layers with tractable Jacobian trace computation. Training requires: (i) samples from the joint distribution of unknown parameters and observed data; and (ii) the posterior score function, which can be computed from physics-based likelihood evaluations, neural operator surrogates, or learned generative models such as conditional normalizing flows. Once trained, the control variates generalize across observations without retraining. We validate our approach on stylized and partial differential equation-constrained Darcy flow inverse problems, outperforming classical Stein control variates and achieving substantial variance reduction, even when the analytical score is replaced by a learned surrogate.

Figures & tables

$d$	Mean Est. VRF	Var. Est. VRF
$2$	$0.009\pm 0.000$	$0.006\pm 0.002$
$4$	$0.017\pm 0.009$	$0.008\pm 0.003$
$8$	$0.034\pm 0.010$	$0.011\pm 0.003$
$16$	$0.083\pm 0.063$	$0.120\pm 0.102$

Table 1: Variance reduction for Gaussian posteriors under per-dimension hyperparameter tuning (recipes in Appendix G ). VRF

=

Var

(h{-}g)

/Var

(h)

; lower is better. Values are mean

\pm

std across

d

components, each averaged over 100 observations.

Figure 1: Dimension scaling of CNCV. (a) VRF for mean and variance estimation across

d\in\{2,4,8,16\}

(lower is better;

\text{VRF}=1

means no improvement). (b) Correlation between

h

and

g

across dimensions.

Figure 2: Sample efficiency for

d=4

mean estimation. (a) VRF is sample-size invariant at

{\sim}\,0.015

. (b) MSE follows

1/N

rate; the CV estimator provides a

{\sim}\,65\times

effective sample increase. Index

j

denotes the component.

Figure 3: Amortized CNCV on Rosenbrock posterior. (a–c) Test observations from left tail, right tail, and ridge (stars) yield diverse posteriors; VRF values are shown in each panel. (d) Per-component VRF for each observation; the same trained model achieves VRF

\in[0.08,0.23]

across all cases (lower is better).

Figure 4: Posterior variance estimation on the Rosenbrock problem. (a) Per-component VRF:

x_{1}

achieves strong reduction while

x_{2}

(banana direction) is harder. (b) Per-observation VRF across 10 test observations (mean

0.44

Figure 5: Corner plot of the nonlinear posterior (

d=4

) for three test observations (blue, red, green). Contours show 50% and 90% HDRs; stars mark the true parameters.

Figure 6: CNCV variance reduction on the nonlinear forward model (

d=4

). (a) Per-component VRF averaged over 10 test observations. (b) Per-observation VRF (mean

0.57

Figure 7: (a) VRF versus total training samples seen. (b) VRF versus ensemble size

L

. Shading shows

\pm 1

standard deviation across test observations. Both panels show Gaussian

d\!=\!4

and Rosenbrock

d\!=\!2

mean estimation, using the reference configuration to isolate training dynamics; the per-dimension-tuned results are in Table 1 .

Problem	Score	VRF	Red. (%)	$\rho$
Gauss. $d\!=\!4$	Analyt.	$0.040\pm 0.011$	96	${>}0.99$
Gauss. $d\!=\!4$	CNF	$0.058\pm 0.002$	94	$0.98$
Rosenb. $d\!=\!2$	Analyt.	$0.08$ – $0.23$	77–92	—
Rosenb. $d\!=\!2$	CNF	$0.010\pm 0.002$	99	${>}0.99$
Nonlin. $d\!=\!4$	Analyt.	$0.57\pm 0.29$	43	—
Nonlin. $d\!=\!4$	CNF	$0.57\pm 0.26$	43	—

Table 2: Analytical vs. CNF-learned score comparison. VRF = Var

(h{-}g)

/Var

(h)

; lower is better. Gaussian and nonlinear: mean

\pm

std across components; Rosenbrock: range over three test observations. Each ensemble is trained and evaluated with the same score source. The Gaussian architecture is held fixed at the reference configuration (

64

hidden units,

3

-layer MLP) to isolate the score-source effect; the per-dimension-tuned results appear in Table 1 .

Figure 8: Sample efficiency with CNF-learned scores. (a) VRF is sample-size invariant. (b) The CV estimator (solid) provides a constant-factor MSE improvement over raw MC (dashed). MSE plateaus at large

N

due to residual bias from the CNF score approximation.

Figure 9: Sample efficiency for Darcy flow. (a) VRF is sample-size invariant at

{\sim}\,0.18

. (b) MSE of the posterior mean estimator follows

1/N

for both raw MC and CNCV, with the CV estimator achieving a

{\sim}\,5.5\times

constant-factor improvement. At large

N

, both curves plateau because the reference mean is estimated from a finite CNF sample pool.

Figure 10: Darcy flow posterior summary. (a) Forward PDE pressure $p^{\dagger}$ with sensor locations (open circles); (b) true log-permeability $u^{\dagger}$ ; (c) CNCV posterior mean from $N=5{,}000$ CNF samples; (d) pointwise posterior standard deviation. All fields reconstructed at $128\times 128$ from KL coefficients.

Figure 11: Variance of the mean estimator on the Darcy flow problem. (a, b) Pixelwise standard deviation of the posterior mean estimator $\hat{\boldsymbol{\mu}}=\frac{1}{N}\sum_{i}\mathbf{x}_{i}$ for vanilla MC and CNCV; (c) std ratio (MC / CNCV), where values ${>}\,1$ indicate that CNCV reduces estimator variance; (d) per-observation VRF across 20 held-out test observations.

Appendix figures & tables18 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 12: Empirical verification of Stein’s identity. Each panel shows the distribution of

\frac{1}{Nd}\sum_{i,j}g_{j}(\mathbf{x}_{i},\mathbf{y})

across 250 test observations (

N\!=\!5{,}000

samples each). All four distributions are centered at zero, confirming the zero-mean guarantee.

Figure 13: Per-KL-component VRF on the Darcy flow problem (

d=100

), sorted by decreasing prior eigenvalue

\lambda_{j}

(orange curve, right axis). High-frequency modes (small

\lambda_{j}

, right side) achieve VRF as low as

0.04

; low-frequency modes (large

\lambda_{j}

, left side) have VRF

\approx 1

. The dashed line marks VRF

=1

(no improvement).

Figure 14: First two KL components of the CNF posterior for three held-out observations. Left column: prior samples (blue) and CNF posterior samples (orange) with the true parameter (red star); the CNF correctly concentrates around the truth. Right column: zoomed view showing the CNF posterior mean (green diamond) near the truth.

Figure 15: Per-component Pearson correlation between the CNF score

\hat{s}_{j}=\partial_{z_{j}}\log p_{\varphi}(\mathbf{z}|\mathbf{y})

and the physics-based posterior score

s_{j}

(adjoint PDE + prior), evaluated at 100 held-out test observations. Mean

\bar{r}=0.94

, median

0.97

Figure 16: Training and validation loss curves. (a) CNF negative log-likelihood: raw per-batch training loss (light), smoothed training loss (dark), and validation NLL per epoch. (b) CNCV mean-squared error: same layout. Both models converge without overfitting.

Figure 17: CNF posterior-mean discrepancy against $20$ held-out truths on the Darcy problem. The signed field-space bias is negligible ( ${\sim}\,-10^{-3}$ ); the unsigned errors—field-space RMSE ( ${\sim}\,0.27$ ) and the per-component KL root-mean-square error, which rises from ${\sim}\,0.28$ on the most data-informed mode toward ${\sim}\,1.0$ on weakly-constrained modes—reflect irreducible posterior spread (Appendix B ) rather than systematic bias.

Figure 18: Darcy posterior summaries for three additional held-out test observations. Each row shows a different observation. Columns: (a) pressure

p^{\dagger}

with sensors, (b) true

u^{\dagger}

, (c) CNCV posterior mean, (d) posterior std.

Panel of Figure 19 — Figure 19: Individual posterior samples for the same three observations as Figure 18 . Each row shows the true field (a) and three independent draws from the CNF approximate posterior (b–d). Samples preserve the frequency content of the truth; the smoothness of the posterior mean in Figure 18 (c) arises from averaging over uncertain high-frequency modes.

Figure 20: Helmholtz FWI posterior summary for a held-out observation. (a) true wave-speed perturbation $\delta v^{\dagger}$ with the $5$ sources (stars) and $40$ receivers (circles); (b) true $\delta v^{\dagger}$ ; (c) CNCV posterior mean; (d) posterior standard deviation.

Figure 21: Per-component VRF on the Helmholtz problem (

d=100

): VRF per KL component (bars), sorted by decreasing prior eigenvalue

\lambda

(orange), averaged over

20

held-out observations. The dashed line marks

\text{VRF}=1

; the data-informed leading modes are hardest, while prior-dominated modes are reduced most.

Figure 22: Sample efficiency on the Helmholtz problem. (a) the mean VRF (over the

d=100

components) is sample-size invariant; (b) the CNCV estimator (solid) attains a constant-factor reduction in MSE over raw Monte Carlo (dashed), both following the

1/N

rate.

	Gauss.	Rosenb.	Nonlin.	Darcy
$d$	2–16	2	4	100
$L$ (ensemble)	16	16	16	16
Tree depth	1–3	2	2	3
MLP layers	5	3	3	5
Hidden units	128	64	64	128
Batch size	2048	2048	2048	4096

Appendix

Table 3: CNCV ensemble hyperparameters.

	Gauss.	Rosenb.	Nonlin.	Darcy
Flow layers	12	12	8	12
Tree depth	2	2	2	6
MLP layers	3	3	3	3
Hidden units	128	128	64	128
Batch size	8192	4096	8192	4096
Train samples	262,144	262,144	262,144	120,000

Appendix

Table 4: Conditional normalizing flow hyperparameters.

Figure 23: Variance reduction on a bimodal Gaussian mixture versus dimension

d

. The degree-one polynomial control variate saturates near

0.4

and does not improve with dimension. The degree-two value at

d=32

(

{\sim}\,2\times 10^{-5}

) is an overfit in-sample estimate, fit without a train/test split, and is unavailable at

d=64

(out of memory). CNCV reaches

\approx 0.023

d=64

, matching the per-observation neural variant while amortizing across observations.

Figure 24: Relative trace error (left axis) and wall-clock time (right axis) of the Hutchinson estimator versus probe count

k

on the Gaussian problem (

d=4

), against the exact Jacobian diagonal.

Figure 25: The same comparison on the Darcy problem (

d=100

); the Hutchinson trace error and its cost relative to the exact diagonal both grow with dimension.

Figure 26: Per-observation VRF under a Student-

t(\nu=5)

likelihood (

d=4

, 20 held-out observations). CNCV reduces variance on every observation (mean

0.10

), whereas the per-observation linear polynomial Stein control variate increases variance on 95% of them (mean

3.5

). The optimal control variate is nonlinear in the saturating Student-

t

score.

Figure 27: VRF versus relative score error on the Gaussian problem (

d=4

), perturbing the analytical score with isotropic Gaussian noise while holding the CNCV ensemble fixed. The VRF degrades gracefully: a 14% error yields

0.032

, and a 57% error still yields

0.35

Variance Reduction

Variance reduction develops methods that make Monte Carlo estimates, gradients, and experimental comparisons more precise with limited data or computation.

See all 107 papers

Oct 7, 2026

Transformed Samplers with Variance Reduction

Markov chain Monte Carlo (MCMC) methods are the standard tool for computing expectations under complex probability distributions. Control variates reduce the variance of the resulting estimates, but a good control variate requires solving the Poisson equation of the sampler, which rarely admits a closed-form solution. Exact solutions are available when the sampler's kernel has a known spectral decomposition on a simple reference density. In our work, we extend these solutions to general targets through a learned change of variables. A bijection, such as a normalizing flow, is trained so that the target becomes close to the reference in a latent space, and we show that Markov kernels and their Poisson solutions are transformed by any bijection. Running such samplers in the latent space then yields explicit control variates, and the estimator is consistent under mild tail conditions on the map and target. Importance sampling (IS) from the flow is the limiting case of the same construction and the control variates apply to it as well. Experiments on synthetic targets and real posteriors compare the procedure against state-of-the-art samplers and control variates.

Siran Liu, Michalis Tisias, Petros Dellaportas

Department of Statistical Science, University College London, UK. · Google DeepMind, UK. · Department of Statistical Science, University College London, UK, and Department of Statistics, Athens University of Economics and Business, Greece.

Markov Chain Monte Carlo Variance Reduction

Ask PDF

May 28, 2026

Deep Adaptive Dimension Reduction for Bayesian Inference in Inverse Problems

Solving high-dimensional PDE-governed inverse problems is often challenging due to complex non-Gaussian posterior distributions, expensive forward model evaluations, and misspecified prior information. To address these issues, we propose a deep adaptive dimension-reduction Bayesian inference framework based on the Variational Flow (VF) model. Since standard normalizing flows are restricted by bijective mappings and cannot directly reduce dimensions, VF overcomes this limitation by integrating VAE-based nonlinear dimension reduction with dual normalizing flows for the latent prior and encoder. This design provides a strictly higher evidence lower bound than VAE and allows more flexible approximation of complex posterior distributions. We further introduce an iterative prior updating strategy that gradually moves the prior mean toward high-probability posterior regions, avoiding manual prior tuning. These components form a closed adaptive loop together with an adaptively fine-tuned Fourier Neural Operator (FNO) surrogate: VF generates posterior-concentrated samples to refine the surrogate, while the updated surrogate further improves posterior inference. Numerical experiments on a 100-dimensional Rosenbrock problem and three standard PDE-governed inverse problems show that our method delivers competitive or superior accuracy compared with MCMC, UKI, and SVGD baselines across all tested configurations, with the most pronounced advantages emerging in challenging scenarios such as high-noise observations and high-dimensional parameter spaces.

Yueyang Wang, Xili Wang, Kejun Tang +3

School of Mathematical Sciences, Peking University · School of Sciences, Great Bay University · Department of Mathematics, Louisiana State University +1

Bayesian Inverse Problems Dimensionality Reduction PDE Inverse Problems

Ask PDF

Mar 14, 2026

Conditional flow matching for physics-constrained inverse problems with finite training data

This study presents a conditional flow matching framework for solving physics-constrained Bayesian inverse problems. In this setting, samples from the joint distribution of inferred variables and measurements are assumed available, while explicit evaluation of the prior and likelihood densities is not required. We derive a simple and self-contained formulation of both the unconditional and conditional flow matching algorithms, tailored specifically to inverse problems. In the conditional setting, a neural network is trained to learn the velocity field of a probability flow ordinary differential equation that transports samples from a chosen source distribution directly to the posterior distribution conditioned on observed measurements. This black-box formulation accommodates nonlinear, high-dimensional, and potentially non-differentiable forward models without restrictive assumptions on the noise model. We further analyze the behavior of the learned velocity field in the regime of finite training data. Under mild architectural assumptions, we show that overtraining can induce degenerate behavior in the generated conditional distributions, including variance collapse and a phenomenon termed selective memorization, wherein generated samples concentrate around training data points associated with similar observations. A simplified theoretical analysis explains this behavior, and numerical experiments confirm it in practice. We demonstrate that standard early-stopping criteria based on monitoring test loss effectively mitigate such degeneracy. The proposed method is evaluated on several physics-based inverse problems. We investigate the impact of different choices of source distributions, including Gaussian and data-informed priors. Across these examples, conditional flow matching accurately captures complex, multimodal posterior distributions while maintaining computational efficiency.

Agnimitra Dasgupta, Ali Fardisi, Mehrnegar Aminy +4

Department of Aerospace & Mechanical Engineering, University of Southern California, Los Angeles, 90089, California, USA · Optimization and Uncertainty Quantification, Sandia National Laboratories, Albuquerque, 87123, New Mexico, USA

Conditional Flow Matching Bayesian Inverse Problems Physics-Informed Generative Modeling Memorization in Generative Models

Ask PDF

Abstract

Variance Reduction

Explore similar work

Transformed Samplers with Variance Reduction

Deep Adaptive Dimension Reduction for Bayesian Inference in Inverse Problems

Conditional flow matching for physics-constrained inverse problems with finite training data