On the Convergence of Gradient Extrapolation Methods for Unbalanced Optimal Transport
share
We study the Unbalanced Optimal Transport (UOT) between two measures of possibly different masses with at most n components, where marginal constraints of the standard Optimal Transport (OT) are relaxed via Kullback-Leibler divergence with regularization factor τ. We propose a novel algorithm based on Gradient Extrapolation Method (GEM-UOT) to find an ε-approximate solution to the UOT problem in O( κ n^2 log(τ n/ε) ), where κ is the condition number depending on only the two input measures. Compared to the only known complexity O(τ n^2 log(n)εlog(log(n)ε)) for solving the UOT problem via the Sinkhorn algorithm, ours is better in ε and lifts Sinkhorn's linear dependence on τ, which hindered its practicality to approximate the standard OT via UOT. Our proof technique is based on a novel dual formulation of the squared ℓ_2-norm regularized UOT objective, which is of independent interest and also leads to a new characterization of approximation error between UOT and OT in terms of both the transportation plan and transport distance. To this end, we further present an algorithm, based on GEM-UOT with fine tuned τ and a post-process projection step, to find an ε-approximate solution to the standard OT problem in O( κ n^2 log( n/ε) ), which is a new complexity in the literature of OT. Extensive experiments on synthetic and real datasets validate our theories and demonstrate the favorable performance of our methods in practice.
READ FULL TEXT