Conditional Optimal Transport (COT) problem aims to find a transport map between conditional source and target distributions while minimizing the transport cost. Recently, these transport maps have been utilized in conditional generative modeling tasks to establish efficient mappings between the distributions. However, classical COT inherits a fundamental limitation of optimal transport, i.e., sensitivity to outliers, which arises from the hard distribution matching constraints. This limitation becomes more pronounced in a conditional setting, where each conditional distribution is estimated from a limited subset of data. To address this, we introduce the Conditional Unbalanced Optimal Transport (CUOT) framework, which relaxes conditional distribution-matching constraints through Csiszár divergence penalties while strictly preserving the conditioning marginals. We establish a rigorous formulation of the CUOT problem and derive its dual and semi-dual formulations. Based on the semi-dual form, we propose Conditional Unbalanced Optimal Transport Maps (CUOTM), an outlier-robust conditional generative model built upon a triangular c-transform parameterization. We theoretically justify the validity of this parameterization by proving that the optimal triangular map satisfies the c-transform relationships. Our experiments on 2D synthetic and image-scale datasets demonstrate that CUOTM achieves superior outlier robustness and competitive distribution-matching performance compared to existing COT-based baselines, while maintaining high sampling efficiency.
Figures & tables
Class
Method
NFE
Checkerboard
Moons
Circles
Swissroll
W2(10−2↓)
W2(10−2↓)
W2(10−2↓)
W2(10−2↓)
Static
PCP-Map [ 3 ]
1
6.27±0.81
8.44±1.09
6.19±0.43
5.35±0.93
COTM
1
7.81±0.66
12.00±0.33
6.34±0.29
8.52±0.32
CUOTM (Ours)
1
6.53±0.40
6.52±0.86
4.46±0.10
4.92±0.22
Dynamic
COT-Flow [ 3 ]
16
8.20±0.49
18.49±2.22
10.04±1.69
6.47±0.69
FM [ 30 ]
596
8.81±0.58
15.55±0.77
7.03±0.17
8.18±0.34
TABLE I: Quantitative results on various 2D synthetic datasets. Baselines are categorized into static (NFE = 1) and dynamic (NFE ≥ 1) frameworks for comparison. Note that COTM corresponds to a special case of CUOTM where Ψ1∗(x)=x and Ψ2∗(x)=x . The results are reported as the average ± standard deviation across five independent trials. Bold and underlined values denote the best and second-best performance, respectively.
Fig. 1: Qualitative results on 2D synthetic datasets. From left to right: target distribution, samples from COTM, samples from CUOTM (Ours), and the respective KDE density visualizations. The top and bottom rows show results for the Moons and Checkerboard datasets, respectively. CUOTM demonstrates superior density recovery and sharper boundary alignment compared to the standard COTM baseline.
Fig. 2: Qualitative results for class-conditional CIFAR-10 generation. Each row corresponds to one of the ten class categories.
Method
NFE
FID ( ↓ )
IFID ( ↓ )
IS ( ↑ )
OT Bayesian Flow [ 10 ]
100
4.10
14.98 †
8.81 †
COTM †
1
33.04
65.66
6.58
CUOTM † (Ours)
1
5.30
15.55
8.79
CUOTM+SD † (Ours)
1
3.71
13.44
8.83
TABLE II: Quantitative results on class-conditional CIFAR-10 comparing conditional OT-based baselines. CUOTM+SD denotes our framework implemented with the α -scheduling [ 35 ] . † indicates the results reproduced by ourselves.
Fig. 3: Qualitative comparison of outlier robustness on the Circles dataset. The in-distribution modes (two center circles) are perturbed with 1% outliers located within annular regions of radius r∈[1.5,2] (top row) and r∈[3,4] (bottom row). While COTM (middle) exhibits significant distortion by attempting to match the outlier, CUOTM (right) effectively ignores the noise by prioritizing high-density regions through marginal relaxation, thereby accurately recovering the majority distribution.
Outlier Range
W2(×10−3,↓)
CUOTM
COTM
[4,5]
0.047
0.205
[3,4]
0.062
0.192
[2,3]
0.070
0.124
[1.5,2]
0.055
0.084
TABLE III: Quantitative comparison of outlier robustness on the Circles dataset. The table reports the Wasserstein-2 distance ( W2×10−3 ) between the majority (clean) distribution and the generated samples under 1% outlier contamination. The outlier range [a,b] denotes that outliers are distributed within an annular region of radius r∈[a,b] from the origin.
Fig. 4: Ablation Study on the cost intensity parameter τ
(Ψ1,Ψ2)
FID( ↓ )
(KL, KL)
3.71
( χ2,χ2 )
5.05
(Softplus, Softplus)
3.79
TABLE IV: Ablation Study on Csiszàr Divergence DΨi(⋅∣⋅) .
Appendix figures & tables2 assets
Supplementary material from the paper’s appendix.
Appendix
Fig. 5: Qualitative results on 2D synthetic datasets. From left to right: target distribution, samples from COTM, samples from CUOTM (Ours), and the respective KDE density visualizations. Each row shows results for the Moons , Checkerboard , Circles , and Swissroll datasets, respectively.
Fig. 6: Qualitative comparison of outlier robustness on the Circles dataset. The in-distribution modes (two center circles) are perturbed with 1% outliers. From top to bottom, the outliers are located within annular regions of radii r∈[1.5,2] , r∈[2,3] , r∈[3,4] , and r∈[4,5] .
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.
Yirong Shen, Mengfei Xia, Junpeng Jing +2
Imperial College London · Ant Group · Brunel University of London
We investigate unpaired image inverse problems, a challenging setting where only independent, non-paired sets of noisy measurements and clean target signals are available for training. We propose a novel inverse problem solver based on Unbalanced Optimal Transport, called Unbalanced Optimal Transport Map for Inverse Problems (UOTIP). Our method formulates the reconstruction task, predicting clean target signals from noisy measurements, as learning a UOT Map from noisy measurement distribution to clean signal distribution by incorporating a likelihood-based cost function. By relaxing the exact marginal constraint, the UOT framework provides key advantages to our model: robustness to multi-level observation noise, adaptability to class imbalance between noisy and clean datasets, and generalizability to diverse noise-type scenarios. Furthermore, we theoretically demonstrate that incorporating a quadratic cost term ensures the existence and uniqueness of the transport map by satisfying the twist condition, even for ill-posed inverse problems. Our experiments demonstrate that UOTIP achieves state-of-the-art performance on unpaired image inverse problem benchmarks, across linear and nonlinear inverse problems.
Donggyu Lee, Taekyung Lee, Jaewoong Choi
Department of Mathematical Sciences, Seoul National University · IPAI (Interdisciplinary Program in Artificial Intelligence, Seoul National University) · Sungkyunkwan University
Current conditioned generative models heavily rely on score functions for guidance during training. When the generative model is a diffusion process with a singular diffusion coefficient and the underlying (conditional) densities either do not exist or are not smooth, we use causal optimal transport to define approximate loss functions that identify a minimum-entropy control for guidance under minimal assumptions. Our approach relies on causal optimal transport and its characterization through the predictable representation property of (conditioned) diffusion processes whose associated martingale problem is well posed, à la Üstünel.
Uğur Aydın, Tamer Başar
Department of Electrical and Computer Engineering and Coordinated Science Laboratory University of Illinois Urbana Champaign Illinois, IL 61801, USA