Optimal Budgeted Rejection Sampling for Generative Models

TL;DR

The paper tackles the problem of generating high-quality samples under a fixed rejection budget in discriminative generative frameworks. It introduces Optimal Budgeted Rejection Sampling (OBRS), derives a closed-form optimal acceptance function that holds for any $f$-divergence, and shows how to compute the budget parameter $c_K$; it further proposes Tw/OBRS to train generators with rejection in mind, yielding flatter loss landscapes and improved mass-covering behavior. Theoretical results characterize precision-recall improvements and provide KL-type bounds via Renyi divergences, while extensive experiments demonstrate superior precision, faster convergence, and enhanced sample diversity across Gaussian benchmarks, BigGAN-CelebA, and diffusion-model settings. The approach offers a practical, principled route to integrate rejection sampling with learning, with potential applicability to broader generative paradigms such as normalizing flows and diffusion models.

Abstract

Rejection sampling methods have recently been proposed to improve the performance of discriminator-based generative models. However, these methods are only optimal under an unlimited sampling budget, and are usually applied to a generator trained independently of the rejection procedure. We first propose an Optimal Budgeted Rejection Sampling (OBRS) scheme that is provably optimal with respect to \textit{any} $f$-divergence between the true distribution and the post-rejection distribution, for a given sampling budget. Second, we propose an end-to-end method that incorporates the sampling scheme into the training procedure to further enhance the model's overall performance. Through experiments and supporting theory, we show that the proposed methods are effective in significantly improving the quality and diversity of the samples.

Figures (20)

• Figure 1: The loss landscape in the parameter domain of a GAN trained on MNIST. The x-axis and y-axis are random directions in the parameter space. The loss is between the target distribution $P$ and the post-rejection distribution. There are three cases: no rejection ($K=1$), $50\%$ acceptance rate ($K=2$) and $20\%$ acceptance rate ($K=5$). OBRS not only reduces loss, but also flattens out the loss landscape and helps avoid local minima.
• Figure 2: Comparing Unbudgeted, DRS azadi_discriminator_2019 and OBRS (ours) for a one-dimensional example. DRS and OBRS are tuned to reach an acceptance ratio of $50\%$. TL: The target and learned distributions $P$ and $\widehat{P}$, along the refined distributions. TR: The acceptance functions for the unbudgeted rejection sampling (dotted black), OBRS (blue), and the DRS (green). BL: The PR-Curves of the different models. BR: Visualisation of the improvements by the OBRS. The straight dotted line corresponds to $\lambda=Kc_K/M$.
• Figure 3: The loss $\mathcal{D}_{\mathrm{GAN}}(P\Vert\widetilde{P})$ is flatten by the OBRS scheme. As the budget $K$ increases, the number of local minima decreases.
• Figure 4: One dimensional example of $\mathcal{D}_f$ minimization: $P$, a mixture of two gaussians is approximated by Gaussian $\widehat{P}=\mathcal{N}\xspace(\mu, \sigma^2)$. The distribution $\widehat{P}$ that minimizes $\mathcal{D}_{\mathrm{GAN}}(P\Vert\widetilde{P})$ leads to a drastically better approximation $\widetilde{P}$ of $P$ than the post rejection distribution induced by the $\widehat{P}$ that minimizes $\mathcal{D}_{\mathrm{GAN}}(P\Vert\widehat{P})$.
• Figure 5: DRS vs. OBRS on a pre-trained BigGAN on CelebA. GAN Baseline model $\widehat{P}_G$, Post-rejection distribution $\widetilde{P}_{a_{\mathrm{DRS}}}$ with DRS, Post-rejection distribution $\widetilde{P}_{a_{\mathrm{O}}}$ with OBRS. (Left) Precision and Recall for different budgets. Lowest budget in black. (Right) FID as a function of the acceptance rate.

Theorems & Definitions (9)

• Theorem 3.1: Optimal Acceptance Function
• Theorem 3.2
• Theorem 3.3: Precision and Recall Improvement
• Lemma 4.1
• Definition A.1: Precision and Recall - sajjadi_assessing_2018
• Definition A.2: Precision and Recall - simon_revisiting_2019
• Corollary B.1
• Theorem C.1
• proof