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

2026-03-07

AI Summary

This work addresses the limitations of traditional conditional optimal transport (COT) in conditional generative modeling, which is sensitive to outliers and constrained by rigid distribution-matching requirements when data subsets are limited. The authors propose a conditional unbalanced optimal transport (CUOT) framework, introducing unbalanced optimal transport into conditional generation for the first time. By relaxing the conditional distribution matching constraint via Csiszár divergences while strictly preserving the conditional marginal distributions, the method incorporates a triangular map parameterization that satisfies the c-transform relationship. Built upon semi-dual optimization, the resulting CUOT model enjoys both theoretical guarantees and enhanced robustness. Experiments demonstrate that CUOT significantly outperforms existing COT approaches on both 2D synthetic and image datasets, exhibiting superior performance in outlier robustness, distribution matching accuracy, and sampling efficiency.

Technology Category

• Constraint Satisfaction and Optimization: Distributed CSP/Optimization
• Machine Learning: Calibration & Uncertainty Quantification
• Computer Vision: Learning & Optimization for CV

Application Category

• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Economics, Online Markets and Human Computation: Trust and reliance of crowd workers and data experts on GenAI

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\'ar 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.

Problem

• Conditional Optimal Transport
• outlier robustness
• conditional generative modeling
• unbalanced optimal transport
• distribution matching

Innovation

• Conditional Unbalanced Optimal Transport
• Outlier Robustness
• c-transform
• Generative Modeling
• Csiszár Divergence