[Paper Review] On Unbalanced Optimal Transport: An Analysis of Sinkhorn Algorithm

Khiem Pham, Khang Le+3 more|arXiv (Cornell University)|Feb 9, 2020

Markov Chains and Monte Carlo MethodsMathematics26 references18 citations

TL;DR

This paper provides a computational complexity analysis of the Sinkhorn algorithm for solving entropic regularized unbalanced optimal transport (UOT), proving a complexity bound of $\widetilde{\mathcal{O}}(n^2 / \varepsilon)$, which is tighter than the best-known bound for standard optimal transport. The analysis leverages geometric convergence of dual updates and scaling properties of the primal solution, demonstrating that UOT enables faster convergence due to relaxed marginal constraints.

ABSTRACT

We provide a computational complexity analysis for the Sinkhorn algorithm that solves the entropic regularized Unbalanced Optimal Transport (UOT) problem between two measures of possibly different masses with at most $n$ components. We show that the complexity of the Sinkhorn algorithm for finding an $\varepsilon$-approximate solution to the UOT problem is of order $\widetilde{\mathcal{O}}(n^2/ \varepsilon)$, which is near-linear time. To the best of our knowledge, this complexity is better than the complexity of the Sinkhorn algorithm for solving the Optimal Transport (OT) problem, which is of order $\widetilde{\mathcal{O}}(n^2/\varepsilon^2)$. Our proof technique is based on the geometric convergence of the Sinkhorn updates to the optimal dual solution of the entropic regularized UOT problem and some properties of the primal solution. It is also different from the proof for the complexity of the Sinkhorn algorithm for approximating the OT problem since the UOT solution does not have to meet the marginal constraints.

Motivation & Objective

• To establish the computational complexity of the Sinkhorn algorithm for solving entropic regularized unbalanced optimal transport (UOT) between discrete measures with possibly different masses.
• To analyze how the absence of marginal constraints in UOT affects convergence speed compared to standard optimal transport (OT).
• To derive a tight upper bound on the number of iterations required to achieve an $\varepsilon$-approximate solution for UOT.
• To validate the theoretical complexity bound through empirical experiments on synthetic data and MNIST.
• To explore the implications of the result for future extensions, such as multi-marginal UOT and UOT barycenter problems.

Proposed method

• Theoretical analysis of the Sinkhorn algorithm's convergence using geometric decay rates of dual variable updates in the entropic regularized UOT problem.
• Derivation of bounds based on scaling properties of the primal solution and the structure of the cost matrix $C$.
• Use of Kullback-Leibler (KL) divergences as penalty functions in the UOT formulation, enabling efficient computation via the Sinkhorn iteration.
• Application of logarithmic and asymptotic approximations to derive a complexity bound with $\widetilde{\mathcal{O}}(n^2 / \varepsilon)$ scaling.
• Empirical validation using synthetic data and MNIST images, comparing theoretical iteration counts ($k_f$) with observed iterations ($k_c$).
• Analysis of convergence ratios $\|v^k - v^*\|_\infty / \|u^{k+1} - u^*\|_\infty$ to assess tightness of theoretical bounds.

Experimental results

Research questions

• RQ1 What is the computational complexity of the Sinkhorn algorithm for approximating unbalanced optimal transport with $\varepsilon$-accuracy?
• RQ2 How does the absence of marginal constraints in UOT affect the convergence rate compared to standard OT?
• RQ3 Can the theoretical complexity bound for UOT be tighter than that of standard OT, and if so, why?
• RQ4 How do the logarithmic terms in the complexity bound relate to the actual convergence behavior in practice?
• RQ5 To what extent can the theoretical iteration count predict empirical performance on real-world datasets like MNIST?

Key findings

• The Sinkhorn algorithm for unbalanced optimal transport achieves a complexity bound of $\widetilde{\mathcal{O}}(n^2 / \varepsilon)$, which is better than the $\widetilde{\mathcal{O}}(n^2 / \varepsilon^2)$ bound for standard optimal transport.
• The improved complexity arises from the relaxation of marginal constraints in UOT, which grants more freedom in updates and accelerates convergence.
• Empirical results on synthetic data show that the ratio $k_f / k_c$ (theoretical to observed iterations) decreases from ~18 to ~8 as $\varepsilon$ decreases from 1 to $10^{-4}$, indicating tighter theoretical bounds in the limit.
• On the MNIST dataset, the empirical log-iteration count $\log(k_c)$ closely follows the theoretical $\log(k_f)$, confirming the validity of the complexity bound.
• The observed convergence ratios of dual variable differences remain close to the theoretical geometric factor $\frac{\tau + \eta}{\tau} = 1.1$, suggesting that the derived inequalities are nearly tight.
• The presence of logarithmic terms in the bound is attributed to suboptimalities in intermediate inequalities, which may be improved with tighter analysis.