Light Unbalanced Optimal Transport

TL;DR

Light Unbalanced OT (U-LightOT) introduces a fast, theoretically grounded solver for continuous unbalanced entropic OT by recasting the problem into a tractable dual-based objective and enforcing a Gaussian mixture parametrization for the UEOT plan. It provides rigorous generalization guarantees and universal approximation under Gaussian mixtures, enabling out-of-sample sampling from the UEOT plan. Empirically, U-LightOT solves UEOT in minutes on CPU and demonstrates robustness to class imbalance in Gaussian mixtures and effective unpaired image translation in FFHQ latent spaces, outperforming several baselines on key metrics while remaining computationally efficient. The approach offers a practical, theoretically supported baseline for UEOT with publicly available code.

Abstract

While the continuous Entropic Optimal Transport (EOT) field has been actively developing in recent years, it became evident that the classic EOT problem is prone to different issues like the sensitivity to outliers and imbalance of classes in the source and target measures. This fact inspired the development of solvers that deal with the unbalanced EOT (UEOT) problem $-$ the generalization of EOT allowing for mitigating the mentioned issues by relaxing the marginal constraints. Surprisingly, it turns out that the existing solvers are either based on heuristic principles or heavy-weighted with complex optimization objectives involving several neural networks. We address this challenge and propose a novel theoretically-justified, lightweight, unbalanced EOT solver. Our advancement consists of developing a novel view on the optimization of the UEOT problem yielding tractable and a non-minimax optimization objective. We show that combined with a light parametrization recently proposed in the field our objective leads to a fast, simple, and effective solver which allows solving the continuous UEOT problem in minutes on CPU. We prove that our solver provides a universal approximation of UEOT solutions and obtain its generalization bounds. We give illustrative examples of the solver's performance. The code is publicly available at https://github.com/milenagazdieva/LightUnbalancedOptimalTransport.

Figures (7)

• Figure 1: Unbalanced EOT problem.
• Figure 2: Conditional plans $\gamma_{\theta,\omega}(y|x)$ learned by our solver in Gaussians Mixture experiment with unbalancedness parameter $\tau\in[10^0,10^1,10^2]$. Here $p_{\omega}$ denotes the normalized first marginal $u_{w}$, i.e., $p_{\omega}=u_{\omega}/\|u_{\omega}\|_1$.
• Figure 3: Visualization of pairs of accuracies (keep-target) for our U-LightOT solver and other OT/EOT methods in the image translation experiment. The values of unbalancedness parameters for our U-LightOT solver ($\tau$) and eyring2023unbalancedness ($\lambda=reg\_m$) are specified directly on the plots.
• Figure 4: Unpaired translation with LightSB, OT-FM, UOT-FM and our U-LightOT solvers applied in the latent space of ALAE pidhorskyi2020adversarial for FFHQ images karras2019style (1024$\times$1024).
• Figure 5: Conditional plans $\gamma_{\theta,\omega}(y|x)$ learned by our solver with scaled $\text{D}_{\chi^2}$ divergences in Gaussians Mixture experiment ($\tau\in[1,2,5,10]$).

Theorems & Definitions (13)

• Theorem 4.1: Tractable form of $\text{D}_{\text{KL}}$ minimization
• Proposition 4.2: Bound for the estimation error
• Theorem 4.3: Gaussian mixture parameterization for the variables provides the universal approximation of UEOT plans
• proof : Proof
• Proposition A.1: Rademacher bound on the estimation error
• proof : Proof of Proposition \ref{['estimation-through-rademacher']}
• Proposition A.2: Bound on the Rademacher complexity of the considered classes
• proof : Proof of Proposition \ref{['proposition-rademacher-complexity']}
• proof : Proof of Proposition \ref{['prop-estimation']}
• Theorem A.3: Fenchel-Rockafellar rockafellar1967duality