Conditional Unbalanced Optimal Transport Maps: An Outlier-Robust Framework for Conditional Generative Modeling

Jiwoo Yoon; Kyumin Choi; Jaewoong Choi

cs.LGMar 7, 2026

Conditional Unbalanced Optimal Transport Maps: An Outlier-Robust Framework for Conditional Generative Modeling

Authors: Jiwoo Yoon, Kyumin Choi, Jaewoong Choi

Organizations: Department of Statistics, Sungkyunkwan University · Department of Applied Artificial Intelligence, Sungkyunkwan University

Abstract

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
Class	Method	NFE	$W_{2}(10^{-2}\!\downarrow)$	$W_{2}(10^{-2}\!\downarrow)$	$W_{2}(10^{-2}\!\downarrow)$	$W_{2}(10^{-2}\!\downarrow)$
Static	PCP-Map [ 3 ]	1	$\underline{6.27\pm 0.81}$	$8.44\pm 1.09$	$6.19\pm 0.43$	$5.35\pm 0.93$
	COTM	1	$7.81\pm 0.66$	$12.00\pm 0.33$	$6.34\pm 0.29$	$8.52\pm 0.32$
	CUOTM (Ours)	1	$6.53\pm 0.40$	$\underline{6.52\pm 0.86}$	$\mathbf{4.46\pm 0.10}$	$\underline{4.92\pm 0.22}$
Dynamic	COT-Flow [ 3 ]	16	$8.20\pm 0.49$	$18.49\pm 2.22$	$10.04\pm 1.69$	$6.47\pm 0.69$
	FM [ 30 ]	596	$8.81\pm 0.58$	$15.55\pm 0.77$	$7.03\pm 0.17$	$8.18\pm 0.34$

TABLE I: Quantitative results on various 2D synthetic datasets. Baselines are categorized into static (NFE = 1) and dynamic (NFE

\geq

1) frameworks for comparison. Note that COTM corresponds to a special case of CUOTM where

\Psi_{1}^{*}(x)=x

and

\Psi_{2}^{*}(x)=x

. The results are reported as the average

\pm

standard deviation across five independent trials. Bold and underlined values denote the best and second-best performance, respectively.

Panel of Figure 2 — 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 ( $\downarrow$ )	IFID ( $\downarrow$ )	IS ( $\uparrow$ )
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

\alpha

-scheduling [ 35 ] . † indicates the results reproduced by ourselves.

Panel of Figure 5 — 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\in[1.5,2]$ (top row) and $r\in[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	$\mathcal{W}_{2}(\times 10^{-3},\downarrow)$
Outlier Range	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 (

\mathcal{W}_{2}\times 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\in[a,b]

from the origin.

Fig. 4: Ablation Study on the cost intensity parameter $\tau$

$(\Psi_{1},\Psi_{2})$	FID( $\downarrow$ )
(KL, KL)	3.71
( $\chi^{2},\chi^{2}$ )	5.05
(Softplus, Softplus)	3.79

TABLE IV: Ablation Study on Csiszàr Divergence

D_{\Psi_{i}}(\cdot|\cdot)

Appendix figures & tables2 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 9 — 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.

Panel of Figure 10 — 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\in[1.5,2]$ , $r\in[2,3]$ , $r\in[3,4]$ , and $r\in[4,5]$ .

Optimal Transport

Optimal transport compares and transforms probability distributions by finding low-cost ways to move mass between them.

See all 289 papers

One-Step Generative Modeling via Unbalanced Optimal Transport
Sep 26, 2026Yirong Shen, Mengfei Xia, Junpeng Jing +2One-Step Generative Modeling Optimal Transport
Imperial College London · Ant Group · Brunel University of London
UOTIP: Unbalanced Optimal Transport Map for Unpaired Inverse Problems
May 20, 2026Donggyu Lee, Taekyung Lee, Jaewoong ChoiImage Reconstruction Optimal Transport
Department of Mathematical Sciences, Seoul National University · IPAI (Interdisciplinary Program in Artificial Intelligence, Seoul National University) · Sungkyunkwan University
Conditioning Degenerate Diffusion Models
Sep 3, 2026Uğur Aydın, Tamer BaşarDiffusion Model Guidance Stochastic Optimal Control
Department of Electrical and Computer Engineering and Coordinated Science Laboratory University of Illinois Urbana Champaign Illinois, IL 61801, USA

Sep 26, 2026J/K move · Enter open · S save

One-Step Generative Modeling via Unbalanced Optimal Transport

Yirong Shen, Mengfei Xia, Junpeng Jing, Lu Gan+1

Imperial College London · Ant Group · Brunel University of London

One-Step Generative Modeling Optimal Transport Wasserstein Gradient Flows

Open paper AskPDF arXiv

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.

Abstract

Optimal Transport

Explore similar work

One-Step Generative Modeling via Unbalanced Optimal Transport

UOTIP: Unbalanced Optimal Transport Map for Unpaired Inverse Problems

Conditioning Degenerate Diffusion Models