🤖 machine learning

Implicit Variational Rejection Sampling

This paper proposes Implicit Variational Rejection Sampling (IVRS), a novel method that combines neural network-based implicit proposal distributions with discriminator-guided rejection sampling and a new IR-ELBO metric to achieve tighter posterior approximations that outperform traditional variational inference techniques.

Original authors: Jian Xu, Shigui Li, Wei Chen, Jiacheng Li, Zhiqi Lin, Delu Zeng, Xinghao Ding, John Paisley, Qibin Zhao

Published 2026-06-15

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Jian Xu, Shigui Li, Wei Chen, Jiacheng Li, Zhiqi Lin, Delu Zeng, Xinghao Ding, John Paisley, Qibin Zhao

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to guess the exact shape of a hidden object in a dark room. You have a flashlight (your computer model) that can shine light on the object, but the beam is a bit fuzzy and doesn't capture every curve perfectly. In the world of machine learning, this "guessing" process is called Variational Inference, and the "hidden object" is a complex mathematical distribution that describes how data is related.

For a long time, scientists used a simple flashlight beam that assumed the object was just a smooth, round ball (this is called "mean-field" approximation). But real-world data is rarely a perfect ball; it's often lumpy, twisted, or has weird holes.

To fix this, researchers started using neural networks to build a smarter, more flexible flashlight that could mold its beam to fit the object's shape better. This is called Implicit Variational Inference. However, even these smart neural networks sometimes get the shape slightly wrong because they are limited by their design or how they were trained.

The New Solution: IVRS

The paper introduces a new method called Implicit Variational Rejection Sampling (IVRS). Think of it as adding a "quality control inspector" to your flashlight setup.

Here is how it works, using a simple analogy:

1. The Proposal (The Guess)
First, your neural network generates a bunch of sample points (like throwing darts) to guess where the hidden object is. Let's call this the "Proposal." Sometimes, these darts land in the right area, but often they land a little too far to the left or right because the neural network isn't perfect.

2. The Inspector (The Discriminator)
This is where the magic happens. The paper adds a second neural network that acts like a strict inspector. This inspector's job is to look at every dart thrown by the first network and ask: "Does this dart land exactly where the true object is, or is it just a guess?"

The inspector doesn't know the exact shape of the object either, but it is very good at comparing the "guess" (the proposal) against the "reality" (the true data). It calculates a score: How likely is this dart to be correct?

3. The Rejection (The Filter)
Now, the system applies a filter:

If the inspector says, "This dart looks great, it's very close to the truth," the system keeps it.
If the inspector says, "This dart is a bit off," the system rejects it and throws it away.

By throwing away the bad guesses and keeping only the best ones, the final collection of darts forms a much more accurate picture of the hidden object than the original neural network could have done alone.

The "Tighter" Bound (IR-ELBO)

The paper also introduces a new mathematical score called IR-ELBO. Imagine you are trying to prove that your flashlight is good. Usually, you have to settle for a "lower bound"—a guarantee that says, "My guess is at least this good."

Because IVRS throws away the bad guesses, the paper proves that the new score (IR-ELBO) is a tighter guarantee. It's like saying, "I don't just know my guess is at least this good; I know it's much closer to the truth because I filtered out the mistakes."

What the Experiments Showed

The authors tested this method on several "toy" problems (simple shapes like bananas or X-shapes) and real-world tasks like:

Predicting house prices or concrete strength (using Bayesian Neural Networks).
Generating images (like handwritten digits from the MNIST dataset or small pictures from CIFAR-10).

In every case, the IVRS method produced results that were closer to the "ground truth" (the real answer) than the traditional methods. It generated sharper images and more accurate predictions.

The Trade-off

There is a small cost to this extra step. Because the system has to "inspect" and potentially "reject" many darts before finding a good one, it takes a little more time and computing power. However, the paper shows that this extra cost is worth it because the final result is significantly more accurate.

In summary: IVRS takes a smart but imperfect guesser (neural network) and pairs it with a strict quality inspector (discriminator). By rejecting the bad guesses, they create a final result that is much closer to reality than what the guesser could achieve on its own.

Technical Summary: Implicit Variational Rejection Sampling (IVRS)

1. Problem Statement

Variational Inference (VI) is a cornerstone of Bayesian machine learning, used to approximate complex posterior distributions. Traditional VI methods often rely on mean-field factorization, which sacrifices posterior expressiveness for computational tractability. To address this, implicit variational inference has emerged, utilizing neural networks to model flexible, implicit distributions without requiring explicit density functions.

However, the paper identifies a critical limitation: despite the theoretical expressiveness of neural networks, practical constraints—such as limited capacity, poor initialization, or optimization difficulties—often prevent these implicit distributions from accurately matching complex true posteriors. Consequently, neural network-based posterior approximations are not always robust "off-the-shelf" solutions. The authors argue that while implicit distributions offer flexibility, they lack the precision required for high-fidelity inference in many scenarios.

2. Methodology: Implicit Variational Rejection Sampling (IVRS)

To bridge the gap between the flexibility of implicit distributions and the precision of exact sampling, the authors propose Implicit Variational Rejection Sampling (IVRS). This method integrates rejection sampling with implicit distributions to refine posterior approximations.

Core Mechanism

The approach treats the implicit distribution generated by a neural network, $q_\phi(z|x)$ , as a proposal distribution. It then applies rejection sampling to generate a resampled distribution, $r_{\theta,\phi}(z|x)$ , which serves as the final variational approximation.

The process involves three key components:

Implicit Proposal: A neural network $f_\phi(x, \epsilon)$ generates samples $z$ from a base noise distribution $\epsilon$ .
Density Ratio Estimation via Adversarial Training: Since the true posterior $p_\theta(z|x)$ $p_{θ} (z ∣ x)$ and the implicit proposal $q_\phi(z|x)$ $q_{ϕ} (z ∣ x)$ often lack analytical density forms, the acceptance probability cannot be computed directly. The authors employ a discriminator network $T_\eta(x, z)$ $T_{η} (x, z)$ trained adversarially to distinguish between samples from the true joint distribution $p(x, z)$ $p (x, z)$ and the proposal distribution.
- The optimal discriminator approximates the log-density ratio: $T^*(x, z) \approx \log q_\phi(z|x) - \log p(z)$ .
- This allows the construction of an acceptance probability function $a(z)$ that is fully differentiable, approximating the ratio between the target (unnormalized posterior) and the proposal.
Rejection Sampling: Samples from the proposal are accepted or rejected based on the estimated acceptance probability. The resulting resampled distribution $r_{\theta,\phi}$ is a refined approximation of the true posterior.

Theoretical Foundation: IR-ELBO

The authors derive a new objective function called the Implicit Resampling Evidence Lower Bound (IR-ELBO).

By applying Jensen's inequality to the marginal log-likelihood using the resampled distribution $r_{\theta,\phi}$ , they derive a tighter variational lower bound than standard VI.
The IR-ELBO incorporates the discriminator network to estimate the implicit density terms, ensuring the objective remains unbiased and suitable for mini-batch optimization.
Monotonicity Property: The paper proves (Proposition 1) that in an idealized setting with an optimal discriminator, the Kullback-Leibler (KL) divergence between the resampled distribution and the true posterior is a monotonically non-increasing function of the scaling parameter $M$ . As $M \to \infty$ , the resampled distribution converges toward the true posterior, theoretically guaranteeing improved accuracy.

3. Key Contributions

The paper outlines three primary contributions:

IVRS Algorithm: The introduction of a novel inference method that combines implicit distributions with rejection sampling to enhance variational inference with neural networks.
IR-ELBO and Theoretical Analysis: The construction of the Implicit Resampling Evidence Lower Bound and a theoretical analysis demonstrating that the resampled distribution reduces KL divergence from the true posterior compared to the raw implicit proposal.
Empirical Validation: Comprehensive experiments showing that IVRS outperforms traditional variational inference techniques (including Adversarial Variational Bayes and Semi-Implicit Variational Inference) in terms of both accuracy and efficiency across various tasks.

4. Experimental Results

The authors evaluated IVRS on several unsupervised learning tasks:

Toy Density Estimation: On low-dimensional synthetic distributions (e.g., Gaussian, Laplace, Banana-shaped, GMM), IVRS consistently reduced cross-entropy values compared to the baseline implicit distributions, demonstrating improved alignment with target contours.
Bayesian Neural Networks (BNNs): On six UCI regression datasets, IVRS achieved competitive or superior performance in terms of Root Mean Squared Error (RMSE) and Negative Log-Likelihood (NLL) compared to AVB, SIVI, and their variants.
Variational Autoencoders (VAEs):
- MNIST: IVRS generated images closer to ground truth than AVB. Quantitatively, it achieved a Negative Log Evidence (NLE) of 81.78, surpassing SIVI and AVB. When integrated into the NVAE architecture, it achieved an NLE of 77.36, improving upon the original NVAE (78.01).
- CIFAR-10 and ImageNet: IVRS achieved lower bits-per-dimension (bpd) scores than classic VAE-based approaches. Notably, adding IVRS to strong baselines like CR-NVAE further reduced bpd (e.g., from 2.51 to 2.43 on CIFAR-10 with augmentation), suggesting the method is complementary to data augmentation.

Efficiency: The authors note that while rejection sampling introduces computational overhead, GPU-accelerated parallel sampling keeps the cost modest (e.g., a 1.82x relative cost increase on Boston dataset at $M=100$ ), which is justified by the significant gains in accuracy.

5. Significance and Claims

The paper positions IVRS as a mechanism to overcome the "expressiveness gap" in neural network-based variational inference. Its significance lies in:

Refinement without Architecture Changes: It improves posterior approximation without necessarily requiring larger or more complex neural network architectures. Ablation studies show that rejection sampling yields greater accuracy gains than simply increasing network capacity (e.g., doubling or tripling hidden units).
Orthogonality to Existing Techniques: The method is presented as complementary to other advancements like data augmentation or specific prior designs (e.g., VampPrior), offering a general inference-refinement step.
Theoretical Guarantee: The derivation of the IR-ELBO and the monotonicity proof provide theoretical support for why rejection sampling tightens the variational bound, moving beyond empirical heuristics.

The authors remain modest regarding limitations, acknowledging that the method requires empirical tuning of the hyperparameter $M$ and that performance in extremely high-dimensional settings (thousands of dimensions) may face challenges due to the curse of dimensionality, though it performs well in moderately high-dimensional scenarios (hundreds of dimensions).