Technical Summary: Conditional Neural Control Variates for Variance Reduction in Bayesian Inverse Problems

Problem Statement

Bayesian inference for inverse problems requires computing expectations under posterior distributions (e.g., posterior means, variances, or predictive quantities). These expectations are typically estimated via Monte Carlo (MC) methods. However, when the quantity of interest $h(x)$ varies significantly under the posterior, accurate estimation demands a prohibitively large number of samples. This cost is particularly acute in partial differential equation (PDE)-constrained inverse problems, where each posterior sample necessitates a computationally expensive forward solve.

While control variates offer a classical remedy by reducing the variance of MC estimators without introducing bias, existing score-based methods (e.g., Stein control variates) suffer from a critical limitation: they require retraining for every new observation $y$. This lack of amortization renders them computationally infeasible when retraining for each observation is too costly. Furthermore, scaling Stein-based control variates to high dimensions is challenging due to the computational cost of computing the divergence of general neural networks, which typically requires $O(d)$ backward passes or stochastic trace estimation.

Methodology

The authors introduce Conditional Neural Control Variates (CNCV), a modular variance reduction method designed to be amortized across observations. The method integrates seamlessly with simulation-based inference pipelines.

1. Conditional Control Variate Construction

The core of the method relies on Stein's identity, which guarantees that for a smooth function $\phi(x, y)$ and a target distribution $p(x|y)$:
$$\mathbb{E}_{p(x|y)}[\nabla_x \cdot \phi(x, y) + \phi(x, y) \cdot \nabla_x \log p(x|y)] = 0$$
This identity defines a family of zero-mean control variates $g(x, y) = \nabla_x \cdot \phi(x, y) + \phi(x, y) \cdot \nabla_x \log p(x|y)$. The authors parameterize $\phi$ using a neural network $\phi_\theta$.

2. Hierarchical Coupling Layer Architecture

To address the computational bottleneck of calculating the divergence $\nabla_x \cdot \phi_\theta$ in high dimensions, the authors propose a specific architecture based on an ensemble of hierarchical affine coupling layers (inspired by HINT).

• Structure: The network organizes affine coupling layers in a recursive binary tree. At each node, the input is split, and the lower partition is transformed based on the upper partition.
• Triangular Jacobian: This structure ensures the Jacobian is lower-triangular, allowing the diagonal elements (and thus the divergence) to be computed exactly in a single forward pass, avoiding the need for $O(d)$ backward passes.
• Ensemble and Permutation: To capture cross-dimension interactions effectively, the method employs an ensemble of $L$ independent networks, each applying a random permutation to the input dimensions. The final control variate is the average of these ensemble members.

3. Training and Inference

• Training Data: The method requires samples from the joint distribution of parameters and data $(x, y)$, which are naturally available in simulation-based inference pipelines (sampling $x$ from the prior and computing $y = F(x) + \epsilon$).
• Score Source: The posterior score $\nabla_x \log p(x|y)$ can be derived from physics-based likelihoods, neural operator surrogates, or learned generative models (e.g., Conditional Normalizing Flows).
• Objective: The network parameters are trained to minimize the mean squared error of the controlled estimator, effectively maximizing the correlation between the control variate and the quantity of interest while penalizing the variance of the variate itself.
• Amortization: Once trained, the CNCV model applies to any new observation $y_{obs}$ without retraining. During inference, posterior samples are drawn (e.g., via a conditional normalizing flow), the score is evaluated, and the learned control variate is subtracted from the raw MC estimate.

Key Contributions

1. Amortized Control Variates: The paper extends neural control variates to the conditional setting, creating a method that integrates with any posterior sampler and score source, eliminating the need for retraining per observation.
2. Tractable High-Dimensional Divergence: The introduction of an ensemble of hierarchical coupling layers with random input permutations enables exact divergence computation in a single forward pass, making the method scalable to high-dimensional problems.
3. Empirical Validation: The method is validated on four inverse problems of increasing complexity (Gaussian, Rosenbrock, nonlinear, and PDE-constrained Darcy flow), demonstrating substantial variance reduction even when the analytical score is replaced by a learned surrogate.

Experimental Results

The authors report results using the Variance Reduction Factor (VRF), where lower values indicate better performance.

• Stylized Problems:
  • Gaussian Posteriors: CNCV achieved VRFs between 0.009 and 0.083 for mean estimation across dimensions $d \in \{2, 4, 8, 16\}$. The method remained robust even when the score was approximated by a Conditional Normalizing Flow (CNF).
  • Rosenbrock (Non-Gaussian): The method successfully generalized across diverse observations (left tail, right tail, ridge) with VRFs ranging from 0.08 to 0.23. It outperformed classical polynomial Stein control variates, which degraded significantly on non-Gaussian targets.
  • Nonlinear Forward Models: Using a denoising diffusion model for sampling and a CNF for the score, CNCV achieved a mean VRF of 0.57.
• PDE-Constrained Darcy Flow:
  • In a 100-dimensional problem involving log-permeability estimation, CNCV achieved a VRF of $\sim 0.18$.
  • This resulted in a $\sim 5.5\times$ reduction in Mean Squared Error (MSE) compared to raw MC.
  • Crucially, because the samples were drawn from a pre-trained CNF, the inference cost was negligible ($\sim 0.5$s) compared to the cost of solving the PDE for raw MC ($\sim 100$s for comparable accuracy).
• Robustness: The method proved robust to score approximation errors. Even with a learned score having a 17.5% relative error, the VRF remained comparable to that achieved with the analytical score.

Significance and Claims

The paper claims that CNCV addresses a fundamental bottleneck in Bayesian inverse problems: the high computational cost of variance reduction in high-dimensional, PDE-constrained settings.

• Amortization: Unlike previous neural control variate methods that require retraining for every observation, CNCV trains once on joint data and generalizes to unseen observations. The authors note that for problems requiring more than roughly twenty distinct observations, the amortized cost of CNCV is preferable to non-amortized baselines.
• Scalability: By leveraging the triangular Jacobian structure of coupling layers, the method overcomes the $O(d)$ computational barrier associated with divergence estimation in high dimensions.
• Flexibility: The method is agnostic to the source of the score, working effectively with both analytical physics-based scores and learned generative models.
• Unbiasedness: The authors emphasize that while the method reduces variance, it does not alter the underlying posterior approximation. If the score and samples come from the same learned source (e.g., a CNF), the estimator remains exactly unbiased for expectations under that learned posterior. Any discrepancy with the true posterior is attributed solely to the quality of the learned score, not the control variate mechanism itself.

The authors conclude that CNCV is particularly valuable in non-Gaussian settings where optimal control variates are nonlinear and amortization is necessary, offering a practical path to efficient uncertainty quantification in complex inverse problems.