Technical Summary: One-Step Generative Modeling via Wasserstein Gradient Flows

Problem Statement

Modern generative models, particularly diffusion models and flow-based methods, have achieved state-of-the-art sample quality but suffer from high computational costs and latency due to their reliance on multi-step iterative sampling. While one-step or few-step distillation methods exist, they often inherit training signals from paths originally designed for iterative generation or rely on heuristic dynamics (e.g., Drifting Models) that lack principled convergence guarantees and are prone to mode collapse. The central challenge is to design a training dynamic that allows a generator to learn a direct, single-step mapping from a reference distribution to a target data distribution while ensuring stable convergence and global mode coverage.

Methodology: W-Flow

The authors propose W-Flow, a framework that decouples the design of training dynamics from the parameterization of the generator. The core idea is to prescribe the evolution of the model distribution toward the target distribution using Wasserstein Gradient Flows (WGF) and then train a static neural network to compress this multi-step evolution into a single inference step.

1. Training Dynamics via Wasserstein Gradient Flow

Instead of learning a sequence of denoising steps or velocity fields directly, W-Flow defines the training trajectory as a continuous-time evolution in the space of probability measures.

• Energy Functional: The evolution is driven by the steepest descent of an energy functional $F(q) = D(q \| p)$, where $q$ is the current model distribution and $p$ is the target data distribution.
• Velocity Field: The continuous-time velocity field $V_t$ is derived from the first variation of the energy functional: $V_t(x) = -\nabla \frac{\delta F}{\delta q}(q_t)(x)$.
• Discretization: The continuous flow is discretized using an explicit Euler scheme. At each training step $k$, the model distribution $q^{(k)}$ is updated via a local transport map $g^{(k)}(x) = x + \eta V^{(k)}(x)$, where $\eta$ is the step size.

2. Instantiation with Sinkhorn Divergence

To ensure the velocity field is computationally tractable and geometrically well-behaved in high dimensions, the authors instantiate the energy functional using the Sinkhorn divergence ($S_\epsilon$), an entropic regularized optimal transport (OT) metric.

• Velocity Field Formulation: The induced velocity field is given by the difference between two barycentric projections:
  $$V^\epsilon_{q,p}(x) = T^\epsilon_{q,p}(x) - T^\epsilon_{q,q}(x)$$
  Here, $T^\epsilon_{q,p}$ transports mass from the generated distribution $q$ to the target $p$, while $T^\epsilon_{q,q}$ represents a self-transport term that enforces consistency with the current distribution.
• Global Coordination: Unlike local kernel interactions (e.g., in MMD or Drifting Models), the Sinkhorn velocity field is derived from a global OT plan, ensuring that particle movements are coordinated across the entire batch to prevent mode collapse.
• Two-Batch Estimator: To avoid self-matching artifacts where particles match to themselves with zero cost, the authors employ a two-batch strategy. They compute the self-transport term $T^\epsilon_{q,q}$ using an independent second batch of samples from the current distribution, rather than masking the diagonal of a single batch's cost matrix.

3. Generator Training

The generator $f_\theta$ is trained to mimic the discrete particle updates defined by the WGF.

• Objective: The loss function minimizes the squared distance between the generator's output and the "target" particle position derived from the Sinkhorn update:
  $$\mathcal{L}(\theta) = \mathbb{E}_{z \sim p_{ref}} \| f_\theta(z) - \text{sg}(f_\theta(z) + \eta V^\epsilon(\hat{q}_\theta, \hat{p})) \|^2$$
  where $\text{sg}(\cdot)$ denotes the stop-gradient operator, treating the updated particle position as a fixed regression target.
• Classifier-Free Guidance (CFG): The framework supports CFG by modifying the velocity field directly. Instead of linearly mixing distributions, the authors propose velocity guidance, which adds a term proportional to the difference between conditional and unconditional velocity fields. This approach is shown to recover the standard exponentially tilted target distribution when using KL divergence and provides a geometrically consistent extension for Sinkhorn divergence.

Key Contributions

1. Principled One-Step Framework: W-Flow establishes a new paradigm for one-step generative modeling by grounding training dynamics in Wasserstein gradient flows, offering a clear variational interpretation absent in heuristic approaches like Drifting Models.
2. Sinkhorn-Based Dynamics: The instantiation of WGF with Sinkhorn divergence provides an efficient, optimal-transport-based update rule that captures global distributional discrepancies and improves mode coverage compared to kernel-based or local interaction methods.
3. Theoretical Convergence: The authors prove that the finite-sample particle dynamics of W-Flow converge to the continuous-time distributional dynamics under suitable regularity assumptions, providing a theoretical justification for the learning process.
4. Two-Batch Estimator: A novel two-batch strategy for estimating self-transport terms is introduced, which eliminates the self-matching bias inherent in single-batch estimators and diagonal masking heuristics.

Experimental Results

The authors evaluate W-Flow on ImageNet 256×256 and FFHQ.

• State-of-the-Art Performance: W-Flow achieves a new state-of-the-art FID of 1.29 (XL scale) and 1.35 (L scale) for one-step class-conditional generation on ImageNet, significantly outperforming previous one-step methods (e.g., Drifting Models, MeanFlow) and surpassing many multi-step diffusion models (e.g., LightningDiT, SiT) despite using only a single function evaluation (NFE).
• Sampling Efficiency: Compared to multi-step diffusion models with similar FID scores, W-Flow yields approximately 100× faster sampling (e.g., 77.88 images/sec vs. 0.93 images/sec for SiT-XL/2 on a single H100 GPU).
• Mode Coverage: In experiments with imbalanced distributions (e.g., 95% senior faces, 5% child faces), W-Flow successfully captures minority modes, whereas Drifting Models exhibit mode collapse, ignoring the minority class entirely.
• Convergence Speed: W-Flow converges substantially faster during training, reaching the performance of a Drifting Model trained for 1280 epochs in only 384 epochs.
• Domain Transfer: The method demonstrates strong capability in domain transfer tasks (e.g., age translation on FFHQ), producing shorter and more coherent transport trajectories compared to heuristic baselines.

Significance and Claims

The paper claims that Wasserstein gradient flows provide a principled and effective foundation for fast and high-fidelity generative modeling. By separating the design of training dynamics from the generator parameterization, W-Flow demonstrates that a single-step generator can be trained to achieve performance comparable to complex multi-step iterative processes. The authors emphasize that the global coordination provided by optimal transport dynamics is key to stabilizing training and ensuring comprehensive mode coverage, addressing critical limitations of existing one-step approaches. The results suggest that W-Flow can serve as a scalable alternative for applications requiring low-latency, high-quality generation, such as real-time creative tools and simulation.