Research map / Probabilistic and scientific modeling / Variational and generative sampling
Scalable Wasserstein Gradient Flow for Generative Modeling through Unbalanced Optimal Transport
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes S-JKO, a scalable Wasserstein gradient flow generative model that uses the semi-dual form of the JKO step (via its equivalence to source-fixed unbalanced OT) plus a reparametrization trick, reducing training complexity from O(K^2) to O(K) and reaching FID 2.62 on CIFAR-10.
Problem
Existing JKO-based Wasserstein gradient flow generative models require training a separate transport map per JKO step, with training complexity O(K^2) in the number of JKO steps, which limits scalability to high-dimensional image data.
Approach
Show the JKO step equals a source-fixed unbalanced OT problem, derive its semi-dual adversarial objective, and reparametrize each step's map as a composition from the source so one-step inference generates samples at any phase, giving O(K) training complexity.
Key results
- S-JKO (Large) reaches FID 2.62 on CIFAR-10, versus 5.55 for NSGF, the second-best WGF-based model.
- S-JKO achieves FID 5.46 on CelebA-HQ-256 and outperforms Source-fixed-UOTM by a large margin.
- With the Small backbone, S-JKO trains in 6 hours with FID 8.78 on CIFAR-10, versus over 30 hours and FID above 20 for prior JKO models.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use multi-phase JKO training instead of one-step Source-fixed-UOTM | S-JKO · CIFAR-10 | FID: 14.4 -> 8.78 | Helped |
| Use multi-phase JKO training instead of one-step Source-fixed-UOTM | S-JKO · CIFAR-10 | FID: 7.53 -> 2.65 | Helped |
| Use multi-phase JKO training instead of one-step Source-fixed-UOTM | S-JKO · CelebA-HQ-256 | FID | Helped |
| Use S-JKO (Large) instead of NSGF | CIFAR-10 · CIFAR-10 | FID: 5.55 -> 2.62 | Helped |
| Vary the number of JKO steps K across K∈{10, 25, 50, 100, 200} (ablation) | S-JKO-KLD · CIFAR-10 | FID: 2.60 | Helped |
| Vary the number of JKO steps K across K∈{10, 25, 50, 100, 200} (ablation) | S-JKO-JSD · CIFAR-10 | FID: 2.66 | Helped |
| Use a very small step size h=0.01 (ablation) | S-JKO with KLD or JSD · CIFAR-10 | FID | Hurt |
| Increase step size to h=0.2 (ablation) | S-JKO-KLD · CIFAR-10 | FID | Hurt |
| Increase step size to h=0.2 (ablation) | S-JKO-JSD · CIFAR-10 | FID | No clear effect |
| Replace KLD with JSD (ablation) | S-JKO · CIFAR-10 and CelebA-HQ | FID | Hurt |
| Replace KLD with JSD (ablation) | S-JKO · CIFAR-10 | robustness to larger h | Helped |
Surprising findings
- Source-fixed UOTM (single-step, K=1) collapses toward the origin on multi-modal 2D targets, while multi-step JKO recovers all modes.
- Performance is robust to the number of JKO steps K, with only marginal gains as K increases.
Limitations
- Source-fixed UOTM and S-JKO use a single fixed Gaussian source; the JKO discretization error grows with large step size h for KLD.
- Relies on adversarial min-max training with R1 regularization, inheriting GAN-style training instability concerns.
Experimental setup
Datasets: CIFAR-10 CelebA-HQ Two Circles 25-Gaussian Mixture
Benchmarks: CIFAR-10 CelebA-HQ-256
Metrics: FID wall-clock training time NFE symmetric KL divergence
Baselines: Source-fixed-UOTM UOTM JKO-ICNN JKO-FlowNet NSGF Sliced-Wasserstein Gradient Flow
Builds on: JKO scheme Source-fixed UOTM NSGF ICNN JKO NCSN++
Closest papers by meaning
- Generalized Wasserstein Flow Matching: Transport Plans, Everywhere, All at Once (NeurIPS 2026 · similarity 0.88)
- Wasserstein Flow Matching: Generative Modeling Over Families of Distributions (ICML 2025 · similarity 0.88)
- Bespoke Solvers for Generative Flow Models (ICLR 2024 · similarity 0.88)
- Flowing Datasets with Wasserstein over Wasserstein Gradient Flows (ICML 2025 · similarity 0.87)
- Wasserstein residuals: Learning Gradient Flows from Population Dynamics (NeurIPS 2026 · similarity 0.86)
- Online Reward-Weighted Fine-Tuning of Flow Matching with Wasserstein Regularization (ICLR 2025 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.