Conditional Neural Control Variates in Bayesian Inference

• Conditional neural control variates are observation-conditioned functions that leverage Stein’s identity and amortized neural networks to achieve unbiased variance reduction in Monte Carlo estimations.
• They employ hierarchical coupling architectures and exact divergence computation to handle high-dimensional, PDE-constrained Bayesian inverse problems efficiently.
• Empirical results demonstrate significant variance reductions (up to 82%) and enhanced sample efficiency, underscoring their practical utility in posterior expectation estimation.

Conditional neural control variates are observation-conditioned control variates for Monte Carlo estimation of posterior expectations in Bayesian inverse problems. In the formulation reported for posterior quantities such as posterior means, variances, or predictive quantities, the objective is to estimate

$I(y)=E_{p(\theta|y)}[h(\theta)]$

under a posterior distribution that may be expensive to sample from and may induce large estimator variance, especially in partial differential equation-constrained settings. The central construction is an amortized function $g(\theta,y)$ with zero posterior mean, so that

$$\hat I_N^{CV} = \frac{1}{N}\sum_{i=1}^N [h(\theta^{(i)})-g(\theta^{(i)},y)]$$

remains unbiased while reducing variance from $\mathrm{Var}[h]/N$ to $\mathrm{Var}[h-g]/N$ (Siahkoohi et al., 24 Feb 2026).

1. Estimation setting and control-variate formulation

The reported setting is Bayesian inference for inverse problems, where posterior expectations are approximated by Monte Carlo. With samples $\theta^{(i)}\sim p(\theta\mid y)$, plain Monte Carlo uses

$$\hat I_N = \frac{1}{N}\sum_{i=1}^N f(\theta^{(i)}).$$

Its variance scales as $\mathrm{Var}(f)/N$, and the method becomes costly when the integrand varies strongly under the posterior or when each sample requires expensive forward solves, such as PDE evaluations (Siahkoohi et al., 24 Feb 2026).

The general control-variate construction starts from a function $c(\theta,y)$ with known posterior expectation $\mu_c(y)=E[c(\theta,y)\mid y]$. The corresponding estimator is

$g(\theta,y)$0

with unbiasedness $g(\theta,y)$1 and variance

$g(\theta,y)$2

Variance is reduced when the control variate is positively correlated with the target and does not introduce excessive variance of its own (Siahkoohi et al., 24 Feb 2026).

Conditional neural control variates specialize this idea by learning an observation-conditioned, amortized zero-mean control variate from joint model-data samples. The conditioning on $g(\theta,y)$3 is central: once trained, the same learned module is evaluated for new observations without retraining. This suggests that the method is designed not merely for variance reduction at a single posterior, but for reuse across an observation family drawn from the same inverse-problem model (Siahkoohi et al., 24 Feb 2026).

2. Stein-based zero-mean construction

The zero-mean property is obtained through Stein’s identity. For any sufficiently smooth $g(\theta,y)$4 satisfying decay at infinity,

$g(\theta,y)$5

Defining the Stein operator as

$g(\theta,y)$6

one obtains a posterior control variate

$g(\theta,y)$7

with exactly zero posterior mean for all $g(\theta,y)$8 (Siahkoohi et al., 24 Feb 2026).

In practice, $g(\theta,y)$9 is parameterized by a neural network and learned so that $$\hat I_N^{CV} = \frac{1}{N}\sum_{i=1}^N [h(\theta^{(i)})-g(\theta^{(i)},y)]$$0 is highly correlated with the target quantity $$\hat I_N^{CV} = \frac{1}{N}\sum_{i=1}^N [h(\theta^{(i)})-g(\theta^{(i)},y)]$$1. The method therefore uses Stein structure to guarantee zero mean and neural parameterization to adapt correlation structure to the quantity of interest. A plausible implication is that the variance-reduction problem is shifted from deriving analytic control variates to learning a Stein-admissible vector field whose induced scalar output tracks posterior fluctuations of $$\hat I_N^{CV} = \frac{1}{N}\sum_{i=1}^N [h(\theta^{(i)})-g(\theta^{(i)},y)]$$2 (Siahkoohi et al., 24 Feb 2026).

The dependence on the posterior score

$$\hat I_N^{CV} = \frac{1}{N}\sum_{i=1}^N [h(\theta^{(i)})-g(\theta^{(i)},y)]$$3

is explicit. The reported implementation allows this score to be computed from physics-based likelihood evaluations, neural operator surrogates, or learned generative models such as conditional normalizing flows; the detailed training procedure also allows analytic differentiation of likelihood plus prior, or approximation by a learned conditional normalizing flow or score-based diffusion model trained on the same joint samples (Siahkoohi et al., 24 Feb 2026).

3. Hierarchical coupling architecture and exact divergence computation

A direct neural parameterization of $$\hat I_N^{CV} = \frac{1}{N}\sum_{i=1}^N [h(\theta^{(i)})-g(\theta^{(i)},y)]$$4 would require $$\hat I_N^{CV} = \frac{1}{N}\sum_{i=1}^N [h(\theta^{(i)})-g(\theta^{(i)},y)]$$5 backward passes to compute the divergence $$\hat I_N^{CV} = \frac{1}{N}\sum_{i=1}^N [h(\theta^{(i)})-g(\theta^{(i)},y)]$$6. To avoid that cost in high dimensions, the method uses an ensemble of $$\hat I_N^{CV} = \frac{1}{N}\sum_{i=1}^N [h(\theta^{(i)})-g(\theta^{(i)},y)]$$7 bijective hierarchical coupling layers with random input permutations $$\hat I_N^{CV} = \frac{1}{N}\sum_{i=1}^N [h(\theta^{(i)})-g(\theta^{(i)},y)]$$8 (Siahkoohi et al., 24 Feb 2026).

For an input split as $$\hat I_N^{CV} = \frac{1}{N}\sum_{i=1}^N [h(\theta^{(i)})-g(\theta^{(i)},y)]$$9, one coupling layer is defined by

$\mathrm{Var}[h]/N$0

where $\mathrm{Var}[h]/N$1 and $\mathrm{Var}[h]/N$2 are small neural networks. The Jacobian is lower triangular, with diagonal entries equal to $\mathrm{Var}[h]/N$3 for $\mathrm{Var}[h]/N$4 and $\mathrm{Var}[h]/N$5 for $\mathrm{Var}[h]/N$6. The corresponding divergence at one node is

$\mathrm{Var}[h]/N$7

By organizing these couplings in a binary tree, every coordinate eventually appears in $\mathrm{Var}[h]/N$8 at some level, and the full Jacobian diagonal, hence the divergence, can be computed exactly in one forward pass via a recursive pass-down of diagonal scalings (Siahkoohi et al., 24 Feb 2026).

An $\mathrm{Var}[h]/N$9-member ensemble with random permutations is used to cover different cross-dimension interactions. The final control variate is the ensemble average

$\mathrm{Var}[h-g]/N$0

The control variate induced by one layer is summarized as

$\mathrm{Var}[h-g]/N$1

The emphasis on exact divergence in one forward pass is important: it is the architectural mechanism by which the Stein construction is made computationally viable in high-dimensional inverse problems (Siahkoohi et al., 24 Feb 2026).

4. Training objective, data requirements, and amortized deployment

Training is posed as minimizing the mean-squared error between the target quantity and the ensemble control variate:

$\mathrm{Var}[h-g]/N$2

Because $\mathrm{Var}[h-g]/N$3 exactly by Stein’s identity, this objective is equivalent up to a constant to minimizing $\mathrm{Var}[h-g]/N$4 (Siahkoohi et al., 24 Feb 2026).

The reported offline training pipeline has four parts. First, one draws $\mathrm{Var}[h-g]/N$5 joint samples from the prior predictive: $\mathrm{Var}[h-g]/N$6 and $\mathrm{Var}[h-g]/N$7. Second, one computes or approximates the posterior score $\mathrm{Var}[h-g]/N$8 either analytically from likelihood plus prior, or with a learned conditional normalizing flow or score-based diffusion model trained on the same joint samples. Third, one evaluates $\mathrm{Var}[h-g]/N$9 by forward passes through the $\theta^{(i)}\sim p(\theta\mid y)$0 coupling networks, using $\theta^{(i)}\sim p(\theta\mid y)$1 in the $\theta^{(i)}\sim p(\theta\mid y)$2 term. Fourth, one updates the network parameters by gradient descent on the sample average of $\theta^{(i)}\sim p(\theta\mid y)$3, with Adam or AdamW (Siahkoohi et al., 24 Feb 2026).

At inference time, the learned module is reused without retraining. For a new observation $\theta^{(i)}\sim p(\theta\mid y)$4, one draws posterior samples $\theta^{(i)}\sim p(\theta\mid y)$5 using any sampler or a conditional normalizing flow, evaluates $\theta^{(i)}\sim p(\theta\mid y)$6, and forms

$\theta^{(i)}\sim p(\theta\mid y)$7

The implementation notes further state that joint training samples come “for free” in simulation-based inference, that the coupling-layer architecture computes divergence exactly in one forward pass, and that once trained, $\theta^{(i)}\sim p(\theta\mid y)$8 is fixed and new observations require no retraining (Siahkoohi et al., 24 Feb 2026).

A common misconception is that observation-conditioned variance reduction must be fit separately for each posterior. The reported formulation rejects that view explicitly: the control variate is amortized over $\theta^{(i)}\sim p(\theta\mid y)$9 and is intended to generalize across observations (Siahkoohi et al., 24 Feb 2026).

5. Empirical performance in stylized and PDE-constrained problems

Empirical results are reported using the variance-reduction factor

$$\hat I_N = \frac{1}{N}\sum_{i=1}^N f(\theta^{(i)}).$$0

with lower values better, $$\hat I_N = \frac{1}{N}\sum_{i=1}^N f(\theta^{(i)}).$$1 perfect, and $$\hat I_N = \frac{1}{N}\sum_{i=1}^N f(\theta^{(i)}).$$2 indicating no gain (Siahkoohi et al., 24 Feb 2026).

Problem setting	Quantity	Reported VRF or gain
Gaussian inverse, $$\hat I_N = \frac{1}{N}\sum_{i=1}^N f(\theta^{(i)}).$$3	Posterior mean	$$\hat I_N = \frac{1}{N}\sum_{i=1}^N f(\theta^{(i)}).$$4–$$\hat I_N = \frac{1}{N}\sum_{i=1}^N f(\theta^{(i)}).$$5
Gaussian inverse, $$\hat I_N = \frac{1}{N}\sum_{i=1}^N f(\theta^{(i)}).$$6	Posterior variance	$$\hat I_N = \frac{1}{N}\sum_{i=1}^N f(\theta^{(i)}).$$7–$$\hat I_N = \frac{1}{N}\sum_{i=1}^N f(\theta^{(i)}).$$8
Rosenbrock inverse, $$\hat I_N = \frac{1}{N}\sum_{i=1}^N f(\theta^{(i)}).$$9	Mean estimation	$\mathrm{Var}(f)/N$0–$\mathrm{Var}(f)/N$1
Nonlinear forward, $\mathrm{Var}(f)/N$2	Mean estimation	$\mathrm{Var}(f)/N$3
Darcy flow, $\mathrm{Var}(f)/N$4 KL modes	Posterior mean	$\mathrm{Var}(f)/N$5

In stylized Gaussian inverse problems with ensemble size $\mathrm{Var}(f)/N$6, posterior-mean estimation achieved $\mathrm{Var}(f)/N$7–$\mathrm{Var}(f)/N$8 across $\mathrm{Var}(f)/N$9, with correlation between $c(\theta,y)$0 and $c(\theta,y)$1 greater than $c(\theta,y)$2. Posterior-variance estimation yielded $c(\theta,y)$3–$c(\theta,y)$4. Sample efficiency was reported as stable as $c(\theta,y)$5 increases, with about $c(\theta,y)$6 effective sample gain at $c(\theta,y)$7 (Siahkoohi et al., 24 Feb 2026).

For the Rosenbrock inverse problem in $c(\theta,y)$8, described as a non-Gaussian banana shape, mean-estimation VRF across three extreme test observations lay in $c(\theta,y)$9. Posterior-variance VRF per dimension was reported as $\mu_c(y)=E[c(\theta,y)\mid y]$0, $\mu_c(y)=E[c(\theta,y)\mid y]$1, with average $\mu_c(y)=E[c(\theta,y)\mid y]$2. In the nonlinear forward problem with $\mu_c(y)=E[c(\theta,y)\mid y]$3 and $\mu_c(y)=E[c(\theta,y)\mid y]$4, mean-estimation VRF was $\mu_c(y)=E[c(\theta,y)\mid y]$5 over ten test observations (Siahkoohi et al., 24 Feb 2026).

The PDE-constrained Darcy flow experiment used $\mu_c(y)=E[c(\theta,y)\mid y]$6 KL modes and $\mu_c(y)=E[c(\theta,y)\mid y]$7 sensors. A $\mu_c(y)=E[c(\theta,y)\mid y]$8M-parameter conditional normalizing flow was trained on $\mu_c(y)=E[c(\theta,y)\mid y]$9k joint $g(\theta,y)$00 samples for both sampling and $g(\theta,y)$01. The CNCV ensemble used $g(\theta,y)$02 and depth $g(\theta,y)$03, yielding $g(\theta,y)$04, described equivalently as $g(\theta,y)$05 variance reduction, and approximately $g(\theta,y)$06 effective sample gain for posterior mean estimation. Per-component VRF was heterogeneous: low-frequency well-constrained modes had $g(\theta,y)$07, high-frequency modes went as low as $g(\theta,y)$08, and the overall average was approximately $g(\theta,y)$09. Posterior-mean maps and standard-deviation fields matched ground truth and reflected sensor geometry (Siahkoohi et al., 24 Feb 2026).

These results support two specific interpretations already encoded in the reported numbers. First, the gains can be substantial but are not uniform across quantities or posterior components. Second, the method does not require low-dimensional posteriors to be effective, since performance is reported in a $g(\theta,y)$10-dimensional PDE-constrained inverse problem (Siahkoohi et al., 24 Feb 2026).

6. Learned-score sensitivity, scope, and relation to earlier neural control variates

The posterior score is a required input to the Stein control variate, but the reported experiments indicate that exact analytical scores are not mandatory. Replacing analytic $g(\theta,y)$11 with conditional-normalizing-flow-learned scores produced nearly identical VRFs: in the Gaussian problem with $g(\theta,y)$12, VRF increased from $g(\theta,y)$13 to $g(\theta,y)$14, and in the nonlinear problem with $g(\theta,y)$15, VRF remained stable at $g(\theta,y)$16 (Siahkoohi et al., 24 Feb 2026). This addresses a second common misconception, namely that Stein-based posterior control variates are practical only when analytical scores are available.

The stated scope is that the method functions as a plug-in, amortized variance-reduction module for any posterior sampler that yields i.i.d. samples and a source of $g(\theta,y)$17. The implementation notes further emphasize physics-based likelihood evaluations, neural operator surrogates, conditional normalizing flows, and diffusion-based score models as admissible score sources (Siahkoohi et al., 24 Feb 2026). A plausible implication is that the method fits naturally within simulation-based inference pipelines in which joint samples and learned posterior surrogates are already present.

In the broader literature on neural control variates, Tucker, Williamson, et al. introduced a general framework for learning low-variance, unbiased gradient estimators for black-box functions of random variables, including an action-conditional extension of advantage actor-critic and the LAX and RELAX estimators (Grathwohl et al., 2017). In that setting, the control variate is a differentiable neural surrogate $g(\theta,y)$18 optimized directly to reduce gradient-estimator variance, rather than a Stein control variate for posterior expectation estimation. The earlier framework therefore belongs to gradient estimation for stochastic optimization, whereas the posterior CNCV construction is targeted at variance reduction of Monte Carlo estimators in Bayesian inverse problems. The shared terminology reflects a common idea—learning conditional neural control variates—but the estimands, unbiasedness arguments, and computational bottlenecks differ (Grathwohl et al., 2017).

Within the reported Bayesian inverse-problem setting, the defining characteristics are thus modularity, amortization over observations, exact zero mean through Stein’s identity, and a coupling-layer architecture that makes divergence evaluation tractable in high dimension (Siahkoohi et al., 24 Feb 2026).