stat.MLFeb 24, 2026
SaveConditional neural control variates for variance reduction in Bayesian inverse problems
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
| Mean Est. VRF | Var. Est. VRF | |
|---|---|---|
Table 1: Variance reduction for Gaussian posteriors under per-dimension hyperparameter tuning (recipes in Appendix G ). VRF Var /Var ; lower is better. Values are mean std across components, each averaged over 100 observations.
Figure 1: Dimension scaling of CNCV. (a) VRF for mean and variance estimation across (lower is better; means no improvement). (b) Correlation between and across dimensions.
Figure 2: Sample efficiency for mean estimation. (a) VRF is sample-size invariant at . (b) MSE follows rate; the CV estimator provides a effective sample increase. Index 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 across all cases (lower is better).
Figure 4: Posterior variance estimation on the Rosenbrock problem. (a) Per-component VRF: achieves strong reduction while (banana direction) is harder. (b) Per-observation VRF across 10 test observations (mean ).
Figure 5: Corner plot of the nonlinear posterior ( ) 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 ( ). (a) Per-component VRF averaged over 10 test observations. (b) Per-observation VRF (mean ).
Figure 7: (a) VRF versus total training samples seen. (b) VRF versus ensemble size . Shading shows standard deviation across test observations. Both panels show Gaussian and Rosenbrock mean estimation, using the reference configuration to isolate training dynamics; the per-dimension-tuned results are in Table 1 .
| Problem | Score | VRF | Red. (%) | |
|---|---|---|---|---|
| Gauss. | Analyt. | 96 | ||
| Gauss. | CNF | 94 | ||
| Rosenb. | Analyt. | – | 77–92 | — |
| Rosenb. | CNF | 99 | ||
| Nonlin. | Analyt. | 43 | — | |
| Nonlin. | CNF | 43 | — |
Table 2: Analytical vs. CNF-learned score comparison. VRF = Var /Var ; lower is better. Gaussian and nonlinear: mean 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 ( hidden units, -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 due to residual bias from the CNF score approximation.
Figure 9: Sample efficiency for Darcy flow. (a) VRF is sample-size invariant at . (b) MSE of the posterior mean estimator follows for both raw MC and CNCV, with the CV estimator achieving a constant-factor improvement. At large , 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 with sensor locations (open circles); (b) true log-permeability ; (c) CNCV posterior mean from CNF samples; (d) pointwise posterior standard deviation. All fields reconstructed at 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 for vanilla MC and CNCV; (c) std ratio (MC / CNCV), where values 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 across 250 test observations ( 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 ( ), sorted by decreasing prior eigenvalue (orange curve, right axis). High-frequency modes (small , right side) achieve VRF as low as ; low-frequency modes (large , left side) have VRF . The dashed line marks VRF (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 and the physics-based posterior score (adjoint PDE + prior), evaluated at 100 held-out test observations. Mean , median .
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 held-out truths on the Darcy problem. The signed field-space bias is negligible ( ); the unsigned errors—field-space RMSE ( ) and the per-component KL root-mean-square error, which rises from on the most data-informed mode toward 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 with sensors, (b) true , (c) CNCV posterior mean, (d) posterior std.
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 with the sources (stars) and receivers (circles); (b) true ; (c) CNCV posterior mean; (d) posterior standard deviation.
Figure 21: Per-component VRF on the Helmholtz problem ( ): VRF per KL component (bars), sorted by decreasing prior eigenvalue (orange), averaged over held-out observations. The dashed line marks ; 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 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 rate.
| Gauss. | Rosenb. | Nonlin. | Darcy | |
| 2–16 | 2 | 4 | 100 | |
| (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 |
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 |
Table 4: Conditional normalizing flow hyperparameters.
Figure 23: Variance reduction on a bimodal Gaussian mixture versus dimension . The degree-one polynomial control variate saturates near and does not improve with dimension. The degree-two value at ( ) is an overfit in-sample estimate, fit without a train/test split, and is unavailable at (out of memory). CNCV reaches at , 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 on the Gaussian problem ( ), against the exact Jacobian diagonal.
Figure 25: The same comparison on the Darcy problem ( ); the Hutchinson trace error and its cost relative to the exact diagonal both grow with dimension.
Figure 26: Per-observation VRF under a Student- likelihood ( , 20 held-out observations). CNCV reduces variance on every observation (mean ), whereas the per-observation linear polynomial Stein control variate increases variance on 95% of them (mean ). The optimal control variate is nonlinear in the saturating Student- score.
Figure 27: VRF versus relative score error on the Gaussian problem ( ), perturbing the analytical score with isotropic Gaussian noise while holding the CNCV ensemble fixed. The VRF degrades gracefully: a 14% error yields , and a 57% error still yields .
More in this topic
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