One-Step Generative Modeling via Wasserstein Gradient Flows

Motivation

• Diffusion and flow methods commonly require many sequential sampling updates, creating substantial latency and computational cost despite their sample quality.The introduction frames the design question as what principles should govern one-step generator training dynamics.Source: §1 · pp. 1–2
• GANs [51] rely on discriminator-induced gradients and minimax optimization, while Drifting Model [14] uses heuristic fields associated with unpredictable convergence and mode collapse.The passage identifies these limitations as motivation for principled one-step training dynamics.Source: §1 · p. 2

W-Flow

• W-Flow evolves the generated distribution by Wasserstein gradient flow using Sinkhorn divergence, then trains a generator to compress those updates into a single inference-time map.The Sinkhorn update combines generated-to-real and generated-to-generated batch-level transport plans.Source: §1 · p. 2

Results

• 1.29 FID at XL scale and 1.35 at L scale make W-Flow state of the art for one-step ImageNet 256×256 generation, surpassing many multi-step diffusion and flow models.Figure 1 plots sample quality against effective sampling compute on ImageNet 256×256.Source: §1 · p. 2, Figure 1 · p. 2

Diffusion and flow-based models

• Diffusion and flow models require many sequential sampling steps; few- and one-step variants often distill teachers or learn shortcuts along a path designed for iterative reverse diffusion.W-Flow instead defines the training-time evolution directly through Wasserstein gradient flows.Source: §2 · pp. 2–3

One-step generator mappings

• GANs [51] [20] use discriminator-driven gradients, while Drifting Model [14] uses a heuristic drifting field; W-Flow focuses on convergence analysis and scaling to ImageNet 256×256.Concurrent work [27] interprets Drifting Model as an approximation to a Sinkhorn-divergence Wasserstein gradient flow.Source: §2 · p. 3

Optimal transport in generative models

• Optimal-transport methods use either OT-inspired training discrepancies or dynamic transport paths, but static OT losses need adversarial or iterative-solver approximations and dynamic methods retain multi-step sampling.The two categories include Wasserstein and Sinkhorn GANs, Schrödinger bridges, and OT-inspired flow models.Source: §2 · p. 3

Wasserstein gradient flows in generative models

• Wasserstein gradient flows describe probability-measure evolution, but JKO-based methods require complex inner optimization at each step, often involving adversarial dual-potential training.These inner problems make JKO-based methods difficult to scale in practice.Source: §2 · p. 3

3.1 W-Flow: Learning one-step transport map under Wasserstein gradient flow

• W-Flow incrementally constructs a transport map, then trains a neural network to realize its local updates while keeping inference to one static generator.The WGF prescribes distributional evolution; the sequence of local transports constructs the one-step map rather than defining a multi-step inference model.Source: §3 · p. 3, §3.1 · pp. 4–5
• The WGF moves the current distribution along the steepest-descent direction of an energy minimized at the target, progressively reducing discrepancy from that target.The velocity field determines intermediate distributions, while composing local maps approximates the flow.Source: §3.1 · p. 4

3.2 Instantiation via Sinkhorn divergence

• The Sinkhorn divergence supplies an efficiently estimable optimal-transport energy whose only stationary distribution is the target, unlike local drifting dynamics that can stall elsewhere.Its velocity balances attraction toward the target with self-interaction through a global mass-conserving coupling, a structure the paper says helps stabilize transport and prevent mode collapse.Source: §3.2 · pp. 5–6

3.3 Interactive particle dynamics

• Using one empirical batch for self-transport causes self-matching bias; the two-batch estimator avoids it, and on a 2D Gaussian it recovers tails missed by the one-batch estimate.The independent batches remove the need for the diagonal-masking heuristic used in [14].Source: §3.3 · p. 6, Figure 3 · p. 6
• Under suitable regularity conditions, the finite-particle, discrete-time dynamics converge to the population continuity-equation dynamics as the step size vanishes and batch sizes grow.The result also states that the population continuity equation has a unique solution over each finite time interval.Source: §3.3 · p. 7

3.4 Implementation

• Training fits the generator to discrete particle updates, using stop-gradient targets to avoid differentiating through the optimal-transport plan and stabilize optimization.The implementation can compute the Sinkhorn divergence in feature space and backpropagate through the feature map.Source: §3.4 · pp. 7–8
• Velocity guidance injects guidance through the difference between conditional and unconditional velocity fields, whereas distribution guidance learns a linearly extrapolated target that deviates from standard CFG.The effective target for velocity guidance is generally unavailable in closed form; with KL energy, the formulation recovers the standard exponentially tilted target.Source: §3.4 · p. 8

4.1 Ablation study

• Sinkhorn performs best among the tested energy functionals; velocity guidance improves FID to 7.08, ε = 0.05 performs best among tested values, and quadratic cost outperforms ℓ2 distance cost.The single-batch self-transport estimate yields 17.57 FID, while the two-batch strategy outperforms the diagonal-mask heuristic.Source: §4.1 · p. 9, Table 2 · p. 9

4.2 Main results on ImageNet

• W-Flow’s FID–Inception Score tradeoff improves over Drifting Model as CFG scale varies, while stronger CFG produces clearer class-specific details and higher perceptual quality.Figure 4 also shows W-Flow L/2 samples as CFG increases from 0.0 to 2.0.Source: §4.2 · pp. 9–10, Figure 4 · p. 10
• FID reaches 1.29 at XL scale and 1.35 at L scale on ImageNet 256×256, establishing a new one-step class-conditional state of the art.W-Flow B/2 also achieves 1.52 FID with 133M parameters, surpassing Drifting Model L/2 with 463M parameters; at 1-NFE it outperforms most diffusion models requiring up to 250 steps.Source: §4.2 · p. 9, Table 3 · p. 10
• Both W-Flow B/2 and L/2 match Drifting Model’s FID after 384 epochs, compared with its 1280-epoch training.The paper attributes faster convergence to training dynamics following the energy functional’s steepest descent.Source: §4.2 · p. 10, Table 4 · p. 10

4.3 Domain transfer and mode coverage

• In domain transfer, W-Flow produces shorter, more coordinated oval-to-circle trajectories and shorter latent transport distances while better preserving source identity in FFHQ age translation.The FFHQ task maps senior faces to young-adult faces in a pretrained autoencoder’s 512-dimensional latent space.Source: §4.3 · pp. 10–11, Figure 5 · p. 11
• W-Flow preserves all six dominant and two minority Gaussian-mixture modes, and captures the 5% child-face mode while maintaining fidelity on the dominant senior-face mode.Drifting Model collapses on dominant clusters and nearly drops the minority child mode.Source: §4.3 · p. 11, Figure 6 · p. 11

5 Conclusion

• W-Flow uses Wasserstein gradient flow and Sinkhorn divergence to guide one-step image generation, and achieves 1.29 FID on ImageNet 256×256.The authors identify broader energy functionals and settings such as text-to-image as future directions.Source: §5 · p. 11

A.1 Wasserstein gradient flows of energy functionals

• The appendix derives Wasserstein gradient-flow velocity fields for MMD, KL divergence, and Sinkhorn divergence, with particle-based estimators for MMD and KL.It connects the resulting estimators to the ablation study.Source: §A.1 · pp. 17–18
• For Sinkhorn divergence, the first variation of the entropic OT cost with respect to the model marginal is given by its optimal dual potential.The appendix uses this variation to derive the Sinkhorn velocity field.Source: §A.1 · pp. 17–18
• The KL velocity field combines attraction toward high-density target regions with repulsion from the current model’s score; estimating it requires score functions or KDE approximations.The appendix assumes smooth positive target and model densities for the KL derivation.Source: §A.1 · pp. 19–20

A.2 More discussion on the estimators for self-transport

• A single-batch self-transport estimate has self-matching bias, while diagonal masking replaces it with a batch-size-sensitive bias.Masking excludes each particle from its own target space.Source: §A.2 · p. 21
• The diagonal-mask bias forces particles toward neighbors even at true density peaks and scales as O(1/B).As B shrinks, the forced redistribution increasingly pushes particles away from local density centers.Source: §A.2 · p. 21
• An independent second batch permits near-identical particle matches without an infinite-cost penalty, consistently estimating the continuous self-transport map even at small batch sizes.The OT plan is computed between two independent empirical measures.Source: §A.2 · p. 21

A.3 More discussions on velocity guidance

• Under KL energy, the velocity-guidance field is a WGF toward the exponentially tilted target, which is the unique global minimizer of the combined energy.The derivation identifies the guidance field with the negative spatial gradient of the energy’s first variation.Source: §A.3 · pp. 22–23
• On the 2D 3-mode Gaussian mixture, increasing velocity guidance concentrates mass in the highest-density regions, while distribution guidance scatters particles away from modes.Figure 8 compares the two guidance strategies on conditional generation.Source: §A.3 · p. 23, Figure 8 · p. 23

A.4 More discussions on related work

• Unlike Drifting Model’s heuristic local attraction and repulsion, W-Flow uses Sinkhorn-induced global OT couplings to coordinate particle movement under marginal constraints.The resulting transport mechanism is described as principled and globally consistent.Source: §A.4 · p. 24
• W-Flow uses explicit Euler updates and Sinkhorn scaling to reduce learning to regression, avoiding JKO inner optimization loops and adversarial training.The authors state that this enables scalable learning of a one-step pushforward map.Source: §A.4 · p. 24
• Unlike Sinkhorn GANs, W-Flow uses the transport plan as an empirical oracle without backpropagating through Sinkhorn iterations, avoiding the associated instability and memory overhead.The authors state this allows the method to scale effectively.Source: §A.4 · p. 24
• W-Flow’s OT-based dynamics and direct one-step map contrast with MMD moment matching’s kernel interactions and prescribed interpolation requiring intermediate simulations.The cited discussion notes MMD’s vanishing-gradient risk when distributions are far apart.Source: §A.4 · pp. 24–25

B Proofs

• The proof setup treats a generic velocity field under Assumption B.1 and verifies the assumption for Sinkhorn velocity when the initial and target distributions have bounded support.The stated verification is on a bounded class of measures, not globally on P2(R^d).Source: B Proofs · p. 25

B.1 Complete statement and the proof of Theorem 3.1

• Assumption B.1 imposes measurability and L2 conditions on the velocity field, together with spatial Lipschitz and Wasserstein stability bounds that support well-posedness and convergence.The proof derives stability and uniqueness, then combines empirical-measure convergence with an Euler discretization estimate to obtain the quantitative bound.Source: §B.1 · pp. 25, 27, 28, 30
• As N, M → ∞ and η → 0, the interpolated empirical measure converges quantitatively to the unique solution of the nonlinear continuity equation.The theorem also establishes existence and uniqueness of that solution under Assumption B.1.Source: §B.1 · p. 26

B.2 Assumption verification for Sinkhorn

• With bounded initial and target supports, the Sinkhorn flow, finite-particle dynamics, and piecewise-linear interpolation remain supported in a bounded ball on every finite time interval.The support bound enables verification of Assumption B.1 along the localized flow and its finite-particle approximations.Source: §B.2 · p. 31
• On the localized bounded domain, compact-domain stability and regularity of Sinkhorn barycentric projections verify Assumption B.1, with constants depending only on the time horizon, support radii, and ε.The spatial estimate is required only within the localized domain.Source: §B.2 · pp. 31–32
• The support statement follows from Proposition B.5, and the remaining proof verifies Assumption B.1 on B̄_R^T.Source: §B.2 · p. 32
• Jensen’s inequality yields the L2-growth estimate, establishing (28) on B̄_R^T.Source: §B.2 · pp. 32–33
• Lemma B.6 supplies the q-stability estimate, which is integrated with respect to eq to obtain (29).Source: §B.2 · p. 33
• For p-stability, the self-interaction term cancels, and integration with respect to q gives (30).Source: §B.2 · p. 33

C Additional implementation details

• ImageNet configurations are listed in Table 5, and the reported L/2 training run takes around three days on the stated H100 setup.The experiments largely follow [14]’s configurations without extensive hyperparameter tuning.Source: C Additional implementation details · p. 33, Table 5 · p. 34
• FFHQ experiments use a 512-dimensional pretrained ALAE latent space, with 60,000 training images and a 4-layer, 1024-hidden-unit MLP.The setup uses batch sizes N = M = 64 and a fixed step size η = 1 for the Sinkhorn divergence.Source: C Additional implementation details · p. 33
• For FFHQ, domain transfer maps senior-face source data to young-adult targets, while the mode-coverage task uses batches containing 95% senior and 5% child faces.The domain-transfer network uses residual identity initialization; unconditional mode coverage starts from a Gaussian prior.Source: C Additional implementation details · p. 33
• The toy energy landscape places training checkpoints along the x-axis and evaluates Sinkhorn divergence over a two-dimensional parameter-space slice.A transverse direction is chosen orthogonal to the checkpoint-trajectory directions.Source: C Additional implementation details · p. 34

• W-Flow yields significant sampling speedup while achieving competitive FID on ImageNet 256×256 against multi-step SiT and LightningDiT.The comparison uses 250 sampling steps for the multi-step models and includes latent preparation and VAE decoding in runtime.Source: D Sampling throughput comparison · p. 35, Table 6 · p. 35
• The throughput comparison benchmarks all models on one Nvidia H100 GPU with batch size 4.Source: D Sampling throughput comparison · p. 35
• Table 6 reports throughput in images per second and computes speedup relative to SiT-XL/2.Source: Table 6 · p. 35

Sinkhorn iterations

• Performance improves as Sinkhorn iterations increase, then plateaus at around 10 iterations.Source: E Additional experimental results · p. 35, Table 7 · p. 35

Velocity guidance

• Velocity guidance improves Sinkhorn- and KL-divergence-driven WGF, with KL plus velocity guidance reaching 8.46 FID and outperforming [14].The paper connects velocity guidance with KL divergence to an exponentially tilted target distribution.Source: E Additional experimental results · p. 35, Table 8 · p. 35

Entropic regularization parameter ε

• With velocity guidance, ℓ2-distance OT cost yields higher FIDs than quadratic OT cost, while both perform best at ε = 0.05.Source: E Additional experimental results · p. 35, Table 9 · p. 35

ImageNet 256×256 samples

• Figures 10–12 show uncurated ImageNet 256×256 W-Flow samples at L/2 and XL/2, with CFG w = 0.15 or 2.0.Source: F Additional generated samples · p. 35, Figure 10 · p. 36, Figure 11 · p. 37
• Figure 11 also shows uncurated XL/2 and B/2 samples with CFG w = 0.15.Source: Figure 11 · p. 37

FFHQ mode-coverage samples

• In the FFHQ mode-coverage experiment, none of the Drifting Model’s 100 randomly generated images is a child’s face.Figures 13 and 14 show samples for the Drifting Model and W-Flow, respectively.Source: F Additional generated samples · p. 36, Figure 13 · p. 39

Limitations

• The evaluation covers ImageNet 256×256 and FFHQ; higher-resolution, text-conditioned, video, and other multimodal settings remain future work.Source: G Limitations & Broader impacts · p. 39
• Large-scale experiments require pretrained feature encoders or autoencoders, and the choice of feature encoders remains heuristic.Source: G Limitations & Broader impacts · p. 39
• The convergence analysis covers particle dynamics under regularity and asymptotic assumptions, not finite-network optimization or finite-compute high-dimensional training.Source: G Limitations & Broader impacts · p. 39

Broader impacts

• The paper identifies risks including impersonation, disinformation, biased representations, and privacy concerns, especially in face-related applications.Source: G Limitations & Broader impacts · p. 39
• In the W-Flow mode-coverage experiment, 4 of 100 generated images are child faces, close to the down-sampled population’s 5% ratio.Source: Figure 14 · p. 40
• For releases beyond research settings, the paper recommends safeguards such as provenance tracking, watermarking, usage restrictions, documentation, and monitoring or filtering.Source: G Limitations & Broader impacts · p. 40