Transformed Samplers and Variance Reduction: Solving the Control Variate Problem with Normalizing Flows
Train a normalizing flow to map complex targets to an analytically tractable space, enabling explicit control variates for MCMC variance reduction.
This paper addresses the core challenge of high MCMC variance and the difficulty of constructing control variates, proposing a solution based on learnable variable transformations. Existing control variate methods rely on closed-form Poisson equation solutions, which are only feasible when the target distribution admits a simple spectral decomposition. The key insight is training a normalizing flow (bijection) to map a complex target distribution near a structurally simple reference density — the Markov kernel and its Poisson equation solution are "transported" along with the bijection, yielding explicit control variates for arbitrary targets. Theoretically, estimators remain consistent under mild tail conditions; practically, importance sampling is identified as a limiting case of the framework, demonstrating its generality.
An Old Problem Haunting Monte Carlo Methods
Markov Chain Monte Carlo (MCMC) methods are the standard tool for computing expectations under complex probability distributions. Whether it's posterior estimation in Bayesian inference or partition function calculations in statistical physics, MCMC is virtually everywhere. But it has a persistent flaw — estimator variance is often very high, meaning a large number of samples are needed to obtain reliable results.
Control variates are a classic technique for reducing this variance. The idea is straightforward: construct an auxiliary variable whose expectation is known and which is highly correlated with the target function, then use it to "cancel out" the random fluctuations in the estimate. The catch is that a truly effective control variate requires solving the Poisson equation of the sampler — and this equation almost never has a closed-form solution.
The Poisson equation plays a central role in Markov chain theory: given a target function $f$ and transition kernel $P$, the Poisson equation takes the form $(I - P)g = f - \pi(f)$, where $\pi(f)$ is the expectation of $f$ under the stationary distribution. The solution $g$ is precisely the desired control variate — adding $g(X_{t+1}) - Pg(X_t)$ to the estimator can dramatically reduce variance without changing the expectation. Intuitively, $g$ captures the "accumulated correlations" of the Markov chain, allowing us to eliminate redundant information between adjacent samples. The difficulty lies in the fact that $(I - P)$ is an operator defined on a function space, and a closed-form solution only exists when the transition kernel $P$ admits an analytically tractable eigenfunction expansion — a condition that is nearly impossible to satisfy for general high-dimensional target distributions.
Limitations of Existing Methods: Only Simple Cases Are Solvable
Exact solutions to the Poisson equation are only obtainable when the sampler's kernel has a known spectral decomposition with respect to some simple reference density. In other words, existing theoretical tools can only handle target distributions that are "clean" enough in structure — such as Gaussian reference densities.
But real-world target distributions are often far more complex. True Bayesian posteriors can be high-dimensional, multimodal, and heavy-tailed, making spectral decomposition intractable and closed-form Poisson solutions impossible. This is precisely the bottleneck that a new arXiv paper (arXiv:2610.10870v1) aims to break through.
Core Idea: Using a Learnable Variable Transformation as a Bridge
The paper's approach can be summarized in one sentence: learn a variable transformation that maps a complex target back to a reference space where analytic solutions are available.
Specifically, the researchers train a bijection — typically implemented as a normalizing flow — so that the target distribution in the latent space approximates the structurally simple reference density. The key theoretical contribution is the authors' proof that both a Markov kernel and the solution to its Poisson equation are correspondingly transformed under any bijection.
This means that Poisson equations previously solvable only under simple reference densities can now be "transported" to arbitrarily complex target distributions via this learned mapping.
Running the Sampler in Latent Space
Once this mapping is in hand, the overall pipeline becomes clear: running the sampler in latent space naturally yields explicit control variates. Moreover, the authors prove that as long as the mapping and target satisfy mild tail conditions, the resulting estimator is consistent — providing a theoretical guarantee of the method's reliability.
Importance Sampling as a Limiting Case
An interesting finding is that Importance Sampling (IS) from the normalizing flow is exactly the limiting case of this construction. This shows that the framework applies not only to MCMC — the same control variates can also be applied to importance sampling, making it quite general.
Normalizing Flows are a class of deep generative models that construct complex probability distributions by stacking invertible, differentiable transformations. Their core mechanism leverages the change-of-variables formula: if $z = T(x)$ is a bijection, then the density of $x$ can be expressed as the product of the density of $z$ (typically a simple Gaussian) and the absolute value of the Jacobian determinant. During training, the model learns to make the distribution of the latent variable $z$ as close as possible to the standard reference distribution, while keeping both forward and inverse transforms efficiently computable. In this paper's setting, normalizing flows are not used to generate new samples — instead, they provide a coordinate transformation that "flattens" the complex target density into the reference density. It is precisely this transformation that makes the otherwise intractable Poisson equation analytically manageable in the new coordinates.
Importance Sampling (IS) is another classical Monte Carlo integration technique: samples are drawn from a tractable proposal distribution $q$, then corrected with weights $w(x) = \pi(x)/q(x)$ to estimate expectations under the target distribution $\pi$. IS can be viewed as a "degenerate" Markov chain — samples are mutually independent, with chain length 1. When the latent-space sampler in this paper's construction degenerates to i.i.d. sampling, the control variate method naturally converges to the form of importance sampling. This connection reveals the unifying nature of the framework: MCMC and IS are not two independent tools, but rather special cases of the same mapping-based variance reduction framework under different assumptions about chain correlation — allowing the same learned normalizing flow to seamlessly serve both types of estimators.
Experimental Validation: Benchmarking Against State-of-the-Art
The paper conducts experiments in two settings: synthetic target distributions and real posterior distributions. Comparisons are made against current state-of-the-art samplers and control variate methods. The experimental design has a clear intent — to demonstrate that this "transformation + control variate" combination delivers real-world variance reduction benefits, rather than remaining purely theoretical.
Significance of This Work
From a methodological perspective, the greatest value of this paper lies in bridging normalizing flows — a class of deep generative models — with classical variance reduction theory. Normalizing flows have shone brightly in generative modeling in recent years, but here they play a different role: not generating samples, but serving as a "bridge" that transports an intractable Poisson equation into a solvable space.
This idea has potential appeal for practical Bayesian computation and scientific simulation: if the same accuracy can be achieved with fewer samples, that translates directly to lower computational cost. Of course, the training overhead of normalizing flows themselves, and whether the tail conditions actually hold for real-world targets, remain factors to weigh when deploying the method in practice.
For researchers following advances in sampling methods, variational inference, and Monte Carlo techniques, this is a technical direction worth tracking.
Related articles
world2vec: Building Mechanistic World Models with Latent Dynamics Learning
world2vec proposes a reward-free method for learning mechanistic world models via latent dynamics, faithfully reproducing systems over long horizons with at least an order of magnitude fewer parameters than DreamerV3, TD-MPC2, and LeWM, and achieving accurate system identification on 71 attractor benchmarks.
A New KL-DRO Paradigm: Adaptive Scenario Generation Around a Learned Density Model
KL-DRO research proposes centering the KL ball on a learned density model instead of the empirical distribution, with adaptive scenario generation guiding sampling toward the worst case. Experiments show lower cost under distribution shift when records are scarce and the model can extrapolate, with comparable in-distribution cost.
ODE Samplers for Discrete Diffusion Models: Exponential Residual Flows Accelerate Generation
Discrete diffusion models lack deterministic ODE sampling? This article introduces Exponential Residual Flows (ERFs), which introduce continuous residual variables per token to construct piecewise ODEs for deterministic sampling and support consistency distillation acceleration. Superior to prior methods on code and text generation tasks.