One-shot Conditional Sampling: MMD meets Nearest Neighbors

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.

In many real-world problems, the objective is to generate a plausible output given an input. For example, one may want to clean a noisy image or turn a low-resolution image into a high-resolution one. More generally, given some context, we want to draw realistic samples that respect the relationship between input and output. This is the problem of conditional sampling. In this work, we develop a framework called CGMMD, which trains a neural network to generate plausible outputs given an input. This framework is simpler than that of GANs, which rely on a delicate balancing act between two competing networks. Instead, we use a single, direct objective based on comparing distributions using nearby data points. This also allows for faster output generation than existing diffusion models, which take multiple steps during the output generation stage. We prove mathematically that our method learns the correct input-output relationship as more training data becomes available. We tested CGMMD on tasks like image denoising and improving image resolution. While our approach may not always beat the best specialized methods in quality, it provides a strong balance between quality and speed, with strong theoretical guarantees. More broadly, our work highlights a promising direction for conditional sampling methods based on distribution matching, and may help motivate further research on fast and theoretically grounded generative models.