One-shot Conditional Sampling: MMD meets Nearest Neighbors

2025-09-29

AI Summary

This paper addresses efficient conditional sampling under partially observed conditional distributions. We propose the CGMMD framework, which—uniquely—integrates Maximum Mean Discrepancy (MMD) with nearest-neighbor regression to formulate a non-adversarial, directly optimizable training objective, enabling theoretically guaranteed one-shot conditional sampling via a single forward pass. Theoretical contributions include deriving a unified generalization bound for nearest-neighbor functions and proving that the generated distribution converges in probability to the true conditional distribution. Empirically, CGMMD achieves state-of-the-art performance on synthetic data and real-world inverse problems—including image denoising and super-resolution—while requiring only O(1) forward computations at test time, drastically reducing inference complexity. The method thus bridges rigorous theoretical foundations with practical deployability.

Technology Category

• Machine Learning: Deep Generative Models & Autoencoders
• Computer Vision: Generative Adversarial Networks (GANs) for Vision
• Search and Optimization: Sampling/Simulation-based Search

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Graph Algorithms and Modeling for the Web: Algorithms and analysis for incomplete, noisy, or partially observed Web-related graphs
• User Modeling, Personalization and Recommendation: On-Device user modeling, personalization, and recommendation

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.

Problem

• Generate samples from unobserved conditional distributions
• Enable one-shot conditional sampling with low complexity
• Achieve convergence to true conditional distribution theoretically

Innovation

• CGMMD enables one-shot conditional sampling
• Adversary-free direct minimization training objective
• Convergence to true conditional distribution proven