Conditional neural control variates for variance reduction in Bayesian inverse problems

TL;DR

This paper introduces conditional neural control variates, a scalable method that learns variance reduction techniques for Monte Carlo estimations in Bayesian inverse problems, significantly improving efficiency especially in high-dimensional PDE-constrained scenarios.

Contribution

The paper proposes a novel modular approach using neural control variates with Stein's identity, enabling scalable variance reduction in Bayesian inverse problems without retraining for new observations.

Findings

Substantial variance reduction demonstrated in PDE-constrained Darcy flow problems.

Method generalizes across different observations without retraining.

Effective even when using learned surrogates for the score function.

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, demonstrating substantial variance reduction, even when the analytical score is replaced by a learned surrogate.