Sinkhorn Divergences for Unbalanced Optimal Transport

Abstract

This paper extends the formulation of Sinkhorn divergences to the unbalanced setting of arbitrary positive measures, providing both theoretical and algorithmic advances. Sinkhorn divergences leverage the entropic regularization of Optimal Transport (OT) to define geometric loss functions. They are differentiable, cheap to compute and do not suffer from the curse of dimensionality, while maintaining the geometric properties of OT, in particular they metrize the weak * convergence. Extending these divergences to the unbalanced setting is of utmost importance since most applications in data sciences require to handle both transportation and creation/destruction of mass. This includes for instance problems as diverse as shape registration in medical imaging, density fitting in statistics, generative modeling in machine learning, and particles flows involving birth/death dynamics. Our first set of contributions is the definition and the theoretical analysis of the unbalanced Sinkhorn divergences. They enjoy the same properties as the balanced divergences (classical OT), which are obtained as a special case. Indeed, we show that they are convex, differentiable and metrize the weak * convergence. Our second set of contributions studies generalized Sinkkhorn iterations, which enable a fast, stable and massively parallelizable algorithm to compute these divergences. We show, under mild assumptions, a linear rate of convergence, independent of the number of samples, i.e. which can cope with arbitrary input measures. We also highlight the versatility of this method, which takes benefit from the latest advances in term of GPU computing, for instance through the KeOps library for fast and scalable kernel operations.

Faster Unbalanced Optimal Transport: Translation invariant Sinkhorn and 1-D Frank-Wolfe

“…Because the map Ψ 2 • Ψ 1 ( f + λ) = Ψ 2 • Ψ 1 ( f ) + λ, one can assume without loss of generality that all iterates ft satisfy ft (x 0 ) = 0 for some x 0 in the support of α. Thus under similar assumptions as in [Séjourné et al, 2019], we get that the iterates lie in a compact set, which yields existence of a fixed point f satisfying Ψ 2 • Ψ 1 ( f ) = f . Defining ḡ = Ψ 1 ( f ) and composing the previous relation by Ψ 1 , we get ḡ = Ψ 1 • Ψ 2 (ḡ ).…”

“…Following [Chizat et al, 2018, Séjourné et al, 2019, Sinkhorn algorithm maximizes the dual problem (3) by an alternate maximization on the two variables. In sharp contrast with balanced OT, UOT Sinkhorn algorithm might converge slowly, even if ε is large.…”

An Homogeneous Unbalanced Regularized Optimal Transport model with applications to Optimal Transport with Boundary

“…Note that setting ϕ = ı c retrieves (2.4) (balanced regularized OT) and setting ε = 0 retrieves (2.7) (unbalanced OT). This model has been deeply studied in [23]. In particular, authors prove that the dual problem (2.9) can be solved by iterating an adapted version of the Sinkhorn algorithm that reads [23,Def.…”

“…Nonetheless, this model appears to be supported by strong theoretical properties, in particular through the introduction of an "unbiased" version called the Sinkhorn divergences [21,8], presented in Section 2.3. Unbalanced and Regularized OT have been mixed together in the works [4,23] in a setting that covers most UOT models (though not directly the OTB one). However, the resulting Unbalanced Regularized OT model (UROT) may fail to be homogeneous, mostly because of the introduction of the (non-linear) term α ⊗ β.…”

“…However, the resulting Unbalanced Regularized OT model (UROT) may fail to be homogeneous, mostly because of the introduction of the (non-linear) term α ⊗ β. In particular, naive adaptations of [23] to introduce an entropic regularization in the OTB model will suffer with heavy inhomogeneity, hindering its use in practice and calling for the development of an entropic regularization term that would preserve homogeneity.…”