One-Step Generative Modeling via Wasserstein Gradient Flows
Abstract
Diffusion models and flow-based methods have achieved strong results in image generation, but often rely on costly iterative sampling. We introduce W-Flow, a framework that compresses a Wasserstein gradient flow into a one-step neural generator. By minimizing the Sinkhorn divergence, W-Flow transports a reference distribution toward the data distribution through efficient optimal-transport updates that capture global distributional discrepancies. We prove that, under suitable assumptions, the finite-sample training dynamics converge to the continuous-time distributional dynamics. Empirically, W-Flow sets a new state of the art in one-step ImageNet generation, with improved mode coverage and domain transfer. We will also discuss extensions using alternative energy functionals and applications to post-training.