Neural Control Variates with Automatic Integration

TL;DR

This work addresses variance reduction in Monte Carlo integration by enabling arbitrary neural networks to act as control variates through an antiderivative network $G_\theta$, where $\frac{\partial}{\partial x}G_\theta(x)=g(x)$. By combining neural spatial integration with change-of-variables to spatial domains and unbiased CV estimators, the approach broadens the expressive power of CVs beyond traditional heuristics. The method is validated on Walk-on-Sphere PDE solvers for 2D Poisson and 3D Laplace problems, demonstrating unbiasedness across architectures (e.g., CatSIREN, ModSIREN, MGC-SIREN) and achieving lower variance than baselines like NF and POLY, with favorable wall-time behavior in many-query settings. Overall, the framework offers a flexible, principled path to variance reduction in complex spatial integrals, with practical impact on PDE solvers and rendering/MCMC applications that rely on Monte Carlo integration.

Abstract

This paper presents a method to leverage arbitrary neural network architecture for control variates. Control variates are crucial in reducing the variance of Monte Carlo integration, but they hinge on finding a function that both correlates with the integrand and has a known analytical integral. Traditional approaches rely on heuristics to choose this function, which might not be expressive enough to correlate well with the integrand. Recent research alleviates this issue by modeling the integrands with a learnable parametric model, such as a neural network. However, the challenge remains in creating an expressive parametric model with a known analytical integral. This paper proposes a novel approach to construct learnable parametric control variates functions from arbitrary neural network architectures. Instead of using a network to approximate the integrand directly, we employ the network to approximate the anti-derivative of the integrand. This allows us to use automatic differentiation to create a function whose integration can be constructed by the antiderivative network. We apply our method to solve partial differential equations using the Walk-on-sphere algorithm. Our results indicate that this approach is unbiased and uses various network architectures to achieve lower variance than other control variate methods.

Figures (7)

• Figure 1: Overview of our method. (a) We first create a diffeomorphic transformation $\Phi$ that maps integration domain to a hyper-cube $[-1,1]^d$. (b) We generalize AutoInt lindell2021autoint to hyper-cube $[-1, 1]^d$ (Sec \ref{['sec:spatial-integral']}). (c) During training, we directly minimize the variance of the estimator using Monte Carlo estimation (Sec \ref{['sec:loss']}).
• Figure 2: Convergence curve of our CV estimator using various randomly initialized networks. This suggests that our method can produce unbiased control variates estimators from arbitrary network architectures.
• Figure 3: Top: Equal sample comparison of solving 2D Poisson equations using control variates integrating over a 2D circle. Bottom: Plotting of Network prediction and integration reference. Our network can fit the integrand tightly. This leads to an estimator that produces the lowest variance.
• Figure 4: Equal Sample Comparison of applying control variates (CV) when solving Poisson equations within a coin domain. Top: apply CV on forcing integral; Bottom: apply CV on both forcing and recursive integral. In both settings, our method achieves the lowest error.
• Figure 5: Equal Sample Comparison. We visualize a 2D slice of the solution to solving Laplace Equations within the Blub shape domain. Our method produces less noisy results and achieves lower error than baseline methods.