Table of Contents
Fetching ...

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

Jiwoo Yoon, Kyumin Choi, Jaewoong Choi

TL;DR

The Conditional Unbalanced Optimal Transport (CUOT) framework is introduced, which relaxes conditional distribution-matching constraints through Csiszar divergence penalties while strictly preserving the conditioning marginals, and an outlier-robust conditional generative model built upon a triangular $c-transform parameterization is proposed.

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.

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

TL;DR

The Conditional Unbalanced Optimal Transport (CUOT) framework is introduced, which relaxes conditional distribution-matching constraints through Csiszar divergence penalties while strictly preserving the conditioning marginals, and an outlier-robust conditional generative model built upon a triangular $c-transform parameterization is proposed.

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 -transform parameterization. We theoretically justify the validity of this parameterization by proving that the optimal triangular map satisfies the -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.
Paper Structure (35 sections, 7 theorems, 59 equations, 5 figures, 4 tables, 1 algorithm)

This paper contains 35 sections, 7 theorems, 59 equations, 5 figures, 4 tables, 1 algorithm.

Key Result

Theorem 3.1

Assume that $\mathcal{Y}, \mathcal{V},$ and $\mathcal{U}$ are compact metric spaces. Let $\mathcal{P}(\cdot)$ be equipped with the topology of weak convergence and $\eta \in \mathcal{P}(\mathcal{Y} \times \mathcal{V})$ and $\nu \in \mathcal{P}(\mathcal{Y} \times \mathcal{U})$ be given source and tar

Figures (5)

  • Figure 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.
  • Figure 2: 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.
  • Figure 3: Ablation Study on the cost intensity parameter $\tau$
  • Figure 4: 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.
  • Figure 5: 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]$.

Theorems & Definitions (14)

  • Theorem 3.1: Existence and Uniqueness of the Minimizer in CUOT
  • Theorem 3.2: Duality for CUOT
  • Theorem 3.3: Validity of Parameterization for CUOT
  • Definition A.1: Csiszár Divergence
  • Definition A.2: Conditional $p$-Wasserstein Distance
  • Theorem B.1: Existence and Uniqueness of the Minimizer in CUOT
  • proof
  • Theorem B.2: Duality for CUOT
  • proof
  • Theorem B.3: Validity of Parameterization for CUOT
  • ...and 4 more