Generative Modeling through the Semi-dual Formulation of Unbalanced Optimal Transport

TL;DR

This work tackles the sensitivity of OT-based generative models to outliers by introducing a generative model (UOTM) grounded in the semi-dual formulation of Unbalanced OT. By relaxing marginal constraints with Csiszár divergences and optimizing a semi-dual objective, UOTM achieves robust, stable training and fast convergence while delivering strong target distribution matching. The authors establish a theoretical upper bound on marginal divergences that scales with the cost parameter and show stability advantages over standard OT formulations, corroborated by empirical results on CIFAR-10 and CelebA-HQ that competitive OT-based baselines. Overall, the approach offers a practical and scalable path for robust generative modeling with improved outlier handling and convergence behavior, with code available for reproducibility.

Abstract

Optimal Transport (OT) problem investigates a transport map that bridges two distributions while minimizing a given cost function. In this regard, OT between tractable prior distribution and data has been utilized for generative modeling tasks. However, OT-based methods are susceptible to outliers and face optimization challenges during training. In this paper, we propose a novel generative model based on the semi-dual formulation of Unbalanced Optimal Transport (UOT). Unlike OT, UOT relaxes the hard constraint on distribution matching. This approach provides better robustness against outliers, stability during training, and faster convergence. We validate these properties empirically through experiments. Moreover, we study the theoretical upper-bound of divergence between distributions in UOT. Our model outperforms existing OT-based generative models, achieving FID scores of 2.97 on CIFAR-10 and 6.36 on CelebA-HQ-256. The code is available at \url{https://github.com/Jae-Moo/UOTM}.

Figures (16)

• Figure 1: Generated Samples from UOTM trained on Left: CIFAR-10 and Right: CelebA-HQ.
• Figure 1: Target Distribution Matching Test. UOTM achieves a better approximation of target distribution $\nu$, i.e., $T_{\#} \mu \approx \nu$. $\dagger$ indicates the results conducted by ourselves.
• Figure 2: Outlier Robustness Test on Toy dataset with 1% outlier. For each subfigure, Left: Comparison of target density $\nu$ and generated density $T_\#\mu$ and Right: Transport map of the trained model $\left(x, T(x)\right)$. While attempting to fit the outlier distribution, OTM generates undesired samples outside $\nu$ and learns the non-optimal transport map. In contrast, UTOM mainly generates in-distribution samples and achieves the optimal transport map. (For better visualization, the y-scale of the density plot is manually adjusted.)
• Figure 3: Outlier Robustness Test on Image dataset (CIFAR-10 + 1% MNIST). Left: OTM exhibits artifacts on both in-distribution and outlier samples. Right: UOTM attains higher-fidelity samples while generating MNIST-like samples more sparingly, around 0.2%. FID scores of CIFAR-10-like samples are 13.82 for OTM and 4.56 for UOTM, proving that UOTM is more robust to outliers.
• Figure 4: Visualization of OT Map $\boldsymbol{\left(x,T(x)\right)}$ trained on clean Toy data. The comparison suggests that Fixed-$\mu$ and UOTM models are closer to the optimal transport map compared to OTM.

Theorems & Definitions (12)

• Remark 3.1: UOT as a Generalization of OT
• Remark 3.2: $\Psi^*$ Candidate
• Theorem 3.3
• Theorem 3.4: semi-dual3
• Theorem A.1: uot1semi-dual1semi-dual3
• proof
• Lemma A.2
• proof
• Theorem A.3
• proof