One-Step Generative Modeling via Unbalanced Optimal Transport
Abstract
Drifting models enable one-step generation by amortizing distribution transport into training, but this efficiency places greater demands on the transport field estimated at each update. In large-scale training, the field is computed from finite mini-batches of generated and real samples, which provide only imperfect approximations to the underlying distributions. Balanced optimal transport enforces exact mass matching within every mini-batch, making the estimated field sensitive to the particular composition of the real-data batch. We find that generated and real samples should be treated asymmetrically: letting the mass assigned to real samples adapt while keeping every generated sample fully transported improves generation across six feature-space metrics in controlled ablations, and is more robust to the relaxation strength than relaxing both marginals simultaneously, which falls below balanced transport under stronger relaxation. Motivated by this observation, we propose Unbalanced Optimal Transport Gradient Flow (UOT-GF), which keeps the generated-sample marginal fixed and relaxes only the real-data marginal. Under identical settings at DiT-B/2 on ImageNet-256, UOT-GF improves Fréchet Inception Distance (FID) from 1.53 to 1.46 over the balanced W-Flow baseline; scaling the same recipe yields 1.34 and 1.22 FID at L/2 and XL/2, the best FID among the one-step models we compare. We further derive the induced UOT transport force, establish a kinetic Vlasov--Fokker--Planck formulation whose overdamped zero-temperature limit recovers the drifting dynamics, and characterize non-target stationary states together with sufficient conditions for convergence.
Related
- Barron Optimal Transport I: Generative Modeling
- A Network Simplex Framework for Discrete Unbalanced Optimal Transport
- Multiscale Supervised Unbalanced Optimal Transport Flow Matching
- Existence, Equivalence, and Convergence of Unbalanced Optimal Transport
- Gaussian-Restricted Barycenters for KL-Unbalanced Optimal Transport: Variational Theory and Fixed-Point Convergence
BibTeX
@misc{shen2026onestep,
title = {{One-Step Generative Modeling via Unbalanced Optimal Transport}},
author = {Yirong Shen and Mengfei Xia and Junpeng Jing and Lu Gan and Cong Ling},
howpublished = {arXiv},
year = {2026},
month = sep,
doi = {10.48550/arxiv.2609.32708},
eprint = {2609.32708},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2609.32708},
}
RIS
TY - UNPB
TI - One-Step Generative Modeling via Unbalanced Optimal Transport
AU - Yirong Shen
AU - Mengfei Xia
AU - Junpeng Jing
AU - Lu Gan
AU - Cong Ling
PY - 2026
DA - 2026/09/26
PB - arXiv
DO - 10.48550/arxiv.2609.32708
UR - https://arxiv.org/abs/2609.32708
AB - Drifting models enable one-step generation by amortizing distribution transport into training, but this efficiency places greater demands on the transport field estimated at each update. In large-scale training, the field is computed from finite mini-batches of generated and real samples, which provide only imperfect approximations to the underlying distributions. Balanced optimal transport enforces exact mass matching within every mini-batch, making the estimated field sensitive to the particular composition of the real-data batch. We find that generated and real samples should be treated asymmetrically: letting the mass assigned to real samples adapt while keeping every generated sample fully transported improves generation across six feature-space metrics in controlled ablations, and is more robust to the relaxation strength than relaxing both marginals simultaneously, which falls below balanced transport under stronger relaxation. Motivated by this observation, we propose Unbalanced Optimal Transport Gradient Flow (UOT-GF), which keeps the generated-sample marginal fixed and relaxes only the real-data marginal. Under identical settings at DiT-B/2 on ImageNet-256, UOT-GF improves Fréchet Inception Distance (FID) from 1.53 to 1.46 over the balanced W-Flow baseline; scaling the same recipe yields 1.34 and 1.22 FID at L/2 and XL/2, the best FID among the one-step models we compare. We further derive the induced UOT transport force, establish a kinetic Vlasov--Fokker--Planck formulation whose overdamped zero-temperature limit recovers the drifting dynamics, and characterize non-target stationary states together with sufficient conditions for convergence.
ER -