On Unbalanced Optimal Transport: Gradient Methods, Sparsity and Approximation Error

TL;DR

This work addresses Unbalanced Optimal Transport with KL-divergence marginals by introducing a gradient-based solver (GEM-UOT) built on a novel dual for squared L2 UOT. GEM-UOT attains an iteration complexity of O((alpha+beta) kappa log(tau n/(epsilon))) with per-iteration cost near O(n^2), and it yields sparse transport plans unlike entropic methods. A complementary result provides a quantitative link between UOT and standard OT, showing that OT can be retrieved from UOT by tuning tau and applying a post-processing projection (GEM-OT), with an explicit bound on the distance gap O(1/tau). The paper also introduces GEM-RUOT for practical relaxed UOT problems and proves that UOT distances converge toward OT distances as tau grows, alongside empirical validation on synthetic and real data (CIFAR-10, Fashion-MNIST) and a color-transfer demonstration. These contributions offer principled gradient-based techniques for sparse, scalable UOT computation and a rigorous pathway to approximate OT via UOT with provable guarantees.

Abstract

We study the Unbalanced Optimal Transport (UOT) between two measures of possibly different masses with at most $n$ components, where the marginal constraints of standard Optimal Transport (OT) are relaxed via Kullback-Leibler divergence with regularization factor $τ$. Although only Sinkhorn-based UOT solvers have been analyzed in the literature with the iteration complexity of ${O}\big(\tfrac{τ\log(n)}{\varepsilon} \log\big(\tfrac{\log(n)}{\varepsilon}\big)\big)$ and per-iteration cost of $O(n^2)$ for achieving the desired error $\varepsilon$, their positively dense output transportation plans strongly hinder the practicality. On the other hand, while being vastly used as heuristics for computing UOT in modern deep learning applications and having shown success in sparse OT problem, gradient methods applied to UOT have not been formally studied. In this paper, we propose a novel algorithm based on Gradient Extrapolation Method (GEM-UOT) to find an $\varepsilon$-approximate solution to the UOT problem in $O\big( κ\log\big(\frac{τn}{\varepsilon}\big) \big)$ iterations with $\widetilde{O}(n^2)$ per-iteration cost, where $κ$ is the condition number depending on only the two input measures. Our proof technique is based on a novel dual formulation of the squared $\ell_2$-norm UOT objective, which fills the lack of sparse UOT literature and also leads to a new characterization of approximation error between UOT and OT. To this end, we further present a novel approach of OT retrieval from UOT, which is based on GEM-UOT with fine tuned $τ$ and a post-process projection step. Extensive experiments on synthetic and real datasets validate our theories and demonstrate the favorable performance of our methods in practice.

Figures (15)

• Figure 1: Comparison in number of iterations (left) and per-iteration cost (right) between GEM-UOT and Sinkhorn on synthetic data for $\epsilon=1 \to 10^{-4}$.
• Figure 2: Scalability of $\tau$ (for $\varepsilon= 10^{-2}$) on synthetic data and CIFAR-10. Sinkhorn scales linearly while GEM-UOT scales logarithmically in $\tau$.
• Figure 3: Wall-clock time comparison between GEM-OT and Sinkhorn on synthetic dataset.
• Figure 4: Primal gap $f-f^*$ of GEM-UOT and Sinkhorn on CIFAR-10.
• Figure 5: Algorithmic dependence on $\eta$ and the corresponding accuracy achieved by the algorithm on CIFAR-10 for $10^6$ iterations.

Theorems & Definitions (31)

• Definition 1
• Remark 2
• Remark 3
• Lemma 4
• Corollary 5
• Remark 6
• Remark 7
• Lemma 8
• Lemma 9
• Lemma 10