One-shot Conditional Sampling: MMD meets Nearest Neighbors

TL;DR

CGMMD advances conditional sampling by learning a conditional generator that matches the true conditional distribution via the ECMMD discrepancy. It replaces adversarial training with a direct ECMMD-based objective, enabling one-shot sampling in a single forward pass and offering strong non-asymptotic guarantees on sampler accuracy and convergence of the induced conditional distribution. The framework rests on kernel embeddings, nearest-neighbor ECMMD estimation, and the noise-outsourcing representation, with rigorous finite-sample bounds and distributional convergence results. Empirically, CGMMD demonstrates competitive performance on synthetic tasks and practical vision tasks (denoising and 4× super-resolution), while offering substantial test-time speed advantages over diffusion-based approaches.

Abstract

How can we generate samples from a conditional distribution that we never fully observe? This question arises across a broad range of applications in both modern machine learning and classical statistics, including image post-processing in computer vision, approximate posterior sampling in simulation-based inference, and conditional distribution modeling in complex data settings. In such settings, compared with unconditional sampling, additional feature information can be leveraged to enable more adaptive and efficient sampling. Building on this, we introduce Conditional Generator using MMD (CGMMD), a novel framework for conditional sampling. Unlike many contemporary approaches, our method frames the training objective as a simple, adversary-free direct minimization problem. A key feature of CGMMD is its ability to produce conditional samples in a single forward pass of the generator, enabling practical one-shot sampling with low test-time complexity. We establish rigorous theoretical bounds on the loss incurred when sampling from the CGMMD sampler, and prove convergence of the estimated distribution to the true conditional distribution. In the process, we also develop a uniform concentration result for nearest-neighbor based functionals, which may be of independent interest. Finally, we show that CGMMD performs competitively on synthetic tasks involving complex conditional densities, as well as on practical applications such as image denoising and image super-resolution.

Figures (14)

• Figure 1: Schematic overview of CGMMD: Given training data $(\bm{Y}_1,\bm{X}_1),\ldots,(\bm{Y}_n,\bm{X}_n)$, the samples $\mathscr{X}_n = \{\bm{X}_1,\ldots, \bm{X}_n\}$ and auxiliary noise $\bm\eta_1,\ldots,\bm\eta_n$ are passed through the generator $\bm{g}$ to produce samples $\bm{g}(\bm\eta_1,\bm{X}_1),\ldots, \bm{g}(\bm\eta_n,\bm{X}_n)$. These outputs are compared with the observed $\bm{Y}_1,\ldots,\bm{Y}_n$ values using a nearest-neighbor $\left(G(\mathscr{X}_n)\right)$ based estimate of the ECMMD discrepancy (see \ref{['eqn:ECMMD_definition']}) between true and generated conditional distributions. Edges are color-coded to highlight the dependence of each section on the corresponding inputs. After training, sampling is immediate: for any new input $\bm{X}$, independently generate new $\bm{\eta}\sim P_{\bm{\eta}}$ , the trained model $\hat{\bm{g}}$ then produces $\hat{\bm{g}}(\bm\eta,\bm{X})$ as the conditional output. Each component is described in greater details in Section \ref{['sec:background']} and Section \ref{['sec:objective_ECMMD']}.
• Figure 2: Comparison of conditional generators on the Helix benchmark at $\bm{X} = 1$.
• Figure 3: Low and high resolution images for MNIST digits $\{0,1,2,3,4\}$.
• Figure 4: Noisy and denoised MNIST digits $\{5,6,7,8,9\}$ at $\sigma = 0.5$.
• Figure 5: CelebHQ denoising using CGMMD at $\sigma = 0.25$.

Theorems & Definitions (19)

• Lemma 2.1: Noise Outsourcing Lemma
• Remark 4.1
• Theorem 4.1: Simpler version of Theorem \ref{['thm:convergence_general']}
• Corollary 4.1
• Theorem 5.1
• Corollary 5.1
• Lemma 5.1
• Lemma 5.2
• Lemma 5.3
• Lemma 5.4