Stabilized Sparse Scaling Algorithms for Entropy Regularized Transport Problems

Introduction

Optimal transport (OT) is a classical optimization problem dating back to the seminal work of Monge and Kantorovich (see monographs for introduction and historical context). The induced Wasserstein distances lift a metric from a ‘base’ space $(X,d)$ to probability measures over $X$. This is a powerful analytical tool, for example to study PDEs as gradient flows in Wasserstein space . With the increase of computational resources, OT has also become a popular numerical tool in image processing, computer vision and machine learning (e.g. ).

Many ideas have been presented to extend Wasserstein distances to general non-negative measures. We refer to and references therein for some context. A transport-type distance for general multi-channel signals is proposed in .

Computational Optimal Transport

To this day, the computational effort to solve OT problems remains the principal bottleneck in many applications. In particular large problems, or even multi-marginal problems, remain challenging both in terms of runtime and memory demand.

For the linear assignment problem and discrete transport problems there are (combinatorial) algorithms based on the finite dimensional linear programming formulation by Kantorovich, such as the Hungarian method , the auction algorithm , the network simplex and more . Typically, they work for (almost) arbitrary cost functions, but do not scale well for large, dense problems. On the other hand, there are more geometric solvers, relying on the polar decomposition , that tend to be more efficient. There is the famous fluid dynamic formulation by Benamou and Brenier , explicit computation of the polar decomposition , semi-discrete solvers , and solvers of the Monge-Ampère equation among many others. However, these only work on very specific cost functions, most notably the squared Euclidean distance. In a compromise between efficiency and flexibility, several discrete coarse-to-fine solvers have been proposed that adaptively select sparse sub-problems .

Entropy Regularization for Optimal Transport

In entropy regularization of the linear assignment problem is considered to allow application of smooth optimization techniques or the Sinkhorn matrix scaling algorithm . For sufficiently small regularization the true optimal assignment can be extracted from the approximate solution. For increased numerical stability, the Sinkhorn algorithm is also reformulated in the log-domain. Similarly, in the Sinkhorn algorithm is applied to solve an entropy regularized approximation of the discrete optimal transport problem. It is demonstrated that for moderate regularization strengths the algorithm is trivial to parallelize, easy to implement on GPUs and fast. Besides, it is shown that moderate regularization can actually be beneficial for classification applications. Regularization also makes the optimization problem more well-behaved (e.g. uniqueness of optimal coupling, optimal objective differentiable as function of marginal distributions), which led to the first practical numerical method for approximate computation of Wasserstein barycenters . Today, this approach is widely used, for instance .

More recently, the Sinkhorn algorithm has been extended to more general transport-type problems, such as multi-marginal problems and direct computation of Wasserstein barycenters , gradient flows and unbalanced transport problems , resulting in a family of Sinkhorn-like diagonal scaling algorithms.

Convergence Speed of Sinkhorn Algorithm

In the convergence rate of the Sinkhorn algorithm is studied for positive kernel matrices, yielding a global linear convergence rate of the marginals in terms of Hilbert’s projective metric. However, applied to entropy regularized optimal transport, the contraction factor tends to one exponentially, as the regularization approaches zero. Thus, running the algorithm with this particular measure of convergence is often not practically feasible. In the local convergence rate of the Sinkhorn algorithm near the solution is examined, based on a linearization of the iterations. This bound is tighter and more accurately describes the behaviour of the algorithm close to convergence. But these estimates do not apply when one starts far from the optimal solution, which is the usual case for small regularization parameters. In a comparison is made between the Sinkhorn algorithm and the auction algorithm. In particular the role of the entropy regularization parameter is related to the slack parameter $\varepsilon$ of the auction algorithm and it is pointed out that convergence of both algorithms becomes slower, as these parameters approach zero (but small parameters are required for good approximate solutions). For the auction algorithm this can provably be remedied by $\varepsilon$-scaling, where the $\varepsilon$ parameter is gradually decreased during optimization. Analogously, it is suggested to gradually decrease entropy regularization during the Sinkhorn algorithm to accelerate optimization. Consequently, in the following we will also refer to the entropy regularization parameter as $\varepsilon$ and to the gradual reduction scheme as $\varepsilon$-scaling. The ideas of are refined in . In particular, the latter proves convergence of a modified algorithm with ‘deformed iterations’ where $\varepsilon$ is gradually decreased during the iterations, similar to $\varepsilon$-scaling. They show that the primal iterate converges to the unregularized solution if the decrease is sufficiently slow. Unfortunately, the number of iterations to reach a given value of $\varepsilon$ increases exponentially, as $\varepsilon$ decreases. Thus it is “mostly interesting from the theoretical point of view” [43, p. 8].

Limitations of Entropic Transport

Despite its considerable merits, there are some fundamental constraints to the naive entropy regularization approach. Entropy introduces some blur in the optimal assignment. While this may sometimes be beneficial (see above), in many applications it is considered a nuisance (e.g. it quickly smears distinct features in gradient flows), and one would like to run the scaling algorithm with as little regularization as possible. However, a standard implementation has some major numerical limitations, becoming increasingly severe as the regularization approaches zero. The diagonal scaling factors diverge in the limit of vanishing regularization, leading to numerical overflow and instabilities. Moreover, the algorithm requires an increasing number of iterations to converge. In practice this can often be remedied by $\varepsilon$-scaling, but its efficiency is not yet well understood theoretically. Therefore, numerically this limit is difficult to reach. In addition, naively storing the dense kernel matrix requires just as much memory as storing the full cost matrix in standard linear programming solvers and multiplications with the kernel matrix become increasingly slow. Thus, effective heuristics to avoid storing of, and multiplication by, the dense kernel matrix have been conceived, such as efficient Gaussian convolutions or approximation by a pre-factored heat kernel . However, these remedies only work for particular (although relevant) problems, and do not solve the issues of blur and diverging scaling factors.

2 Contribution and Outline

In Section 2 we recall the framework for transport-type problems and corresponding scaling algorithms for their entropy regularized counterparts, as put forward in . The main contributions of this article are twofold: In Section 3 we propose to combine four modifications of the Sinkhorn algorithm to address issues with numerical instability, slow convergence and large kernel matrices. In Section 4 a new convergence analysis for the Sinkhorn algorithm is derived, based on an analogy to the auction algorithm. The two sections can be read independently from each other. The modifications used in Section 3 are:

Section 3.1: A log-domain stabilization of the Sinkhorn algorithm, as described in . It allows to numerically run the algorithm at small regularizations while largely retaining the simple matrix scaling structure.

Section 3.2: The well-known $\varepsilon$-scaling heuristic, to reduce the number of required iterations.

Section 3.3: Sparsification of the kernel matrix by adaptive truncation, to reduce memory demand and accelerate iterations. We quantify the error induced by truncation and propose a truncation scheme which reliably yields small error bounds that are easy to evaluate. While truncation has been proposed elsewhere (e.g. ), to the best of our knowledge the present article gives the first concrete bounds for the inflicted error.

Section 3.4: A multi-scale scheme, inspired, for instance, by . This serves two purposes: First, it allows for a more efficient computation of the truncated kernel. Second, we propose to combine the coarse-to-fine approach with simultaneous $\varepsilon$-scaling, which drastically reduces the number of variables during early stages of $\varepsilon$-scaling, without losing significant precision.

We emphasize that each modification builds on the previous ones (see Remark 5.1) and only combining all four leads to an algorithm that can solve large problems with significantly less runtime, memory and regularization, as compared to the naive algorithm. The adaptations extend to the more general scaling algorithms for transport-type problems presented in .

In Section 4 we develop a new convergence analysis of the Sinkhorn algorithm, based on analogy to the auction algorithm, different from the Hilbert metric approach of . The structure of Section 4 is:

Section 4.1: The classical auction algorithm for the linear assignment problem is recalled.

Section 4.2: A slightly modified asymmetric variant of the Sinkhorn algorithm is given and a bound is derived for the number of iterations until a prescribed accuracy is reached. As for the auction algorithm, for fixed $\varepsilon$ the maximal number of iterations scales as $\mathcal{O}(1/\varepsilon)$. This is in good agreement with numerical experiments (cf. Section 5.2). To avoid the difficulties with slow convergence in Hilbert’s projective metric (cf. Section 1.1) we choose a weaker, but reasonable, measure of convergence (cf. Remark 4.13).

Section 4.3: We prove stability of optimal dual solutions of entropy regularized OT under changes of the regularization parameter. This also implies stability of dual solutions in the limit of vanishing regularization and therefore complements results of (see also Remark 4.17).

Section 4.4: Our eventual goal is a better theoretical understanding of the $\varepsilon$-scaling heuristic and its efficiency. We show that the above stability result is an important step and discuss missing steps for a full proof. To our knowledge (with the exception of , see above), these are the first theoretical results towards $\varepsilon$-scaling for the Sinkhorn algorithm.

Numerical experiments confirm the efficiency of the modified algorithm (Section 5.2). Examples for unbalanced optimal transport, barycenters, and Wasserstein gradient flows illustrate that the modified algorithms retain the versatility of the diagonal scaling algorithms presented in (Section 5.3).

3 Notation and Preliminaries

We assume that the reader has a basic knowledge of convex optimization, such as convex conjugation, Fenchel–Rockafellar duality and primal-dual gaps (cf. ).

For Sect. 4 we require the following Lemma.

The first line follows immediately from $0\leq\exp(a(z)/\varepsilon)\leq\exp(\max a/\varepsilon)$. Line three then follows from $\min(a-b)\leq\max(a)-\max(b)\leq\max(a-b)$. The second and fourth line are implied by $\operatorname*{softmin}(a,\varepsilon)=-\operatorname*{softmax}(-a,\varepsilon)$.

Entropy Regularized Transport-Type Problems and Diagonal Scaling Algorithms

Recently it has been proposed to replace the constraints $\textnormal{P}_{X}\pi=\mu$ and $\textnormal{P}_{Y}\pi=\nu$ by soft constraints. This allows meaningful comparison between measures of different total mass. Such formulations were studied e.g. in (see also for more context). A particularly relevant choice for the soft constraints is the Kullback–Leibler divergence. A corresponding ‘unbalanced’ transport problem is given by

where $\lambda>0$ is a weighting parameter. Note that neither $\mu$, $\nu$ nor $\pi$ need to be probability measures in this case and each may have different total mass.

one obtains the Wasserstein–Fisher–Rao (WFR) distance (or Hellinger–Kantorovich distance), introduced independently and simultaneously in . WFR is the length distance induced by GHK .

Problems (2.1) and (2.2) share a common structure: in both we optimize over non-negative measures $\pi$ on the product space $X\times Y$, there is a linear cost term $\left\langle c,\pi\right\rangle$ and two functions act on the marginals of $\pi$. They are prototypical examples of a family of transport-type optimization problems with a common functional structure that was introduced in . The general structure is given in the following definition.

The indicator function $\iota_{\leq c}(\textnormal{P}_{X}^{\top}\,\alpha+\textnormal{P}_{Y}^{\top}\,\beta)$ denotes the classical optimal transport dual constraint $\alpha(x)+\beta(y)\leq c(x,y)$ for all $(x,y)\in X\times Y$ (see Section 1.3).

This family also covers Wasserstein gradient flows and the structure can be extended to multiple couplings to describe barycenter and multi-marginal problems (see for details). As indicated, the standard optimal transport problem (2.1) is obtained as a special case.

Problem (2.1) is a special case of Def. 2.1 with $F_{X}:=\iota_{\{\mu\}}$ and $F_{Y}:=\iota_{\{\nu\}}$. The primal and dual functional are given by:

Likewise, we can proceed for the unbalanced transport problem (2.2).

Problem (2.2) is a special case of Def. 2.1 with $F_{X}:=\lambda\cdot\operatorname{KL}(\cdot|\mu)$ and $F_{Y}:=\lambda\cdot\operatorname{KL}(\cdot|\nu)$. The primal and dual functional are given by:

2 Entropy Regularization and Diagonal Scaling Algorithms

with the convention $\exp(-\infty)=0$. $K$ is called the kernel associated with $c$ and the regularization parameter $\varepsilon$. For convenience we formally introduce the function

We obtain the regularized equivalent to Def. 2.1.

Primal optimizers $\pi^{\dagger}$ have the form

where $(\alpha^{\dagger},\beta^{\dagger})$ are dual optimizers. Conversely, for dual optimizers $(\alpha^{\dagger},\beta^{\dagger})$, $\pi^{\dagger}$ constructed as above is primal optimal .

Intuitively we see the relation between (2.4) and (2.10) as $\varepsilon\to 0$. For example, the term \varepsilon\,\operatorname{KL}^{\ast}\big{(}[\textnormal{P}_{X}^{\top}\alpha+\textnormal{P}_{Y}^{\top}\beta]/\varepsilon\big{|}K\big{)} in (2.10b) can be interpreted as a smooth barrier function for the dual constraint $\textnormal{P}_{X}^{\top}\alpha+\textnormal{P}_{Y}^{\top}\beta\leq c$ in (2.4b). We refer to Sect. 1.1 for references to rigorous convergence results.

Under suitable assumptions problem (2.10b) can be solved by alternating optimization in $\alpha$ and $\beta$ (see for details). For fixed $\beta$, consider the $\operatorname{KL}^{\ast}$-term:

Note that the last term is constant w.r.t. $\alpha$. Therefore, optimizing (2.10b) over $\alpha$, for fixed $\beta$ corresponds to maximizing

This is a proximal step of $F_{X}$ for the $\operatorname{KL}$ divergence with step size $1/\varepsilon$ (see Def. 1.2). So, by using the PD-optimality conditions between (2.12) and (2.13) (see e.g. [4, Thm. 19.1]), for a given $\beta$ the primal optimizer $\sigma^{\dagger}$ of (2.13) and the dual optimizer $\alpha^{\dagger}$ of (2.12) are given by

Analogously, optimization w.r.t. $\beta$ for fixed $\alpha$ is related to $\operatorname{KL}$ proximal steps of $F_{Y}$. Starting from some initial $\beta^{(0)}$, we can iterate alternating optimization to obtain a sequence $\beta^{(0)},\alpha^{(1)},\beta^{(1)},\alpha^{(2)},\ldots$ as follows:

The algorithm becomes somewhat simpler when it is formulated in terms of the effective variables

For more convenient notation we introduce the proxdiv operator of a function $F$ and step size $1/\varepsilon$:

The primal-dual relation (2.11) then becomes $\pi^{\dagger}=\operatorname{diag}(u^{\dagger})\,K\,\operatorname{diag}(v^{\dagger})$, which is why $u$ and $v$ are often referred to as diagonal scaling factors.

Throughout this article, we will refer to the arguments of the dual functionals (2.4b) and (2.10b) as dual variables and denote them with $(\alpha,\beta)$. The effective, exponentiated variables, introduced in (2.16), will be denoted by $(u,v)$ and referred to as scaling factors.

For future reference let us state the full scaling algorithm.

The stopping criterion is typically a bound on the primal-dual gap between dual iterates $(\alpha,\beta)=\varepsilon\,\log(u,v)$ and primal iterate $\pi=\operatorname{diag}(u)\,K\,\operatorname{diag}(v)$, an error bound on the marginals of $\pi$ (for standard optimal transport) or a pre-determined number of iterations.

With alternating iterations (2.15) or (2.18) a large family of functionals of form (2.10a) can be optimized, as long as the $\operatorname{KL}$ proximal steps of $F_{X}$ and $F_{Y}$ can be computed efficiently. A particularly relevant sub-family is, where $F_{X}$ and $F_{Y}$ are separable and are a sum of pointwise functions. Then the $\operatorname{KL}$ steps decompose into pointwise one-dimensional $\operatorname{KL}$ steps, see [14, Section 3.4] for details.

Since Section 4 focusses on the special case of entropy regularized optimal transport, let us explicitly state the corresponding functional and iterations.

The proximal steps of $F_{X}$ and $F_{Y}$ are trivial (if $K$ has non-empty columns and rows) and we recover the famous Sinkhorn iterations:

Stabilized Sparse Multi-Scale Algorithm

Throughout this section we combine four adaptions to the Algorithm 1 to overcome the limitations of a naive implementation outlined in Section 1.1.

When running Algorithm 1 with small regularization parameter $\varepsilon$, entries in the kernel $K$, and the scaling factors $u$ and $v$ may become both very small and very large, leading to numerical difficulties. However, under suitable conditions (e.g. standard optimal transport, finite cost function) it can be shown that the optimal dual variables $(\alpha,\beta)$ remain finite and have a stable limit as $\varepsilon\to 0$ (, see also Remark 4.17). In and others it was proposed to formulate the Sinkhorn iterations directly in terms of the dual variables, instead of the scaling factors. For example, an update of $\alpha$ would be performed as follows:

As an alternative, we employ the redundant parametrization of the iterations as proposed in. The scaling factors $(u,v)$, (2.16), are written as

Analogous to the function $\textnormal{get}K$, (2.9), we define the stabilized kernel as

The second line, (3.3b), should be used for numerical evaluation such that extreme values in $(\alpha,\beta)$ and $c$ can cancel before exponentiation. Moreover, we introduce a stabilized version of the proxdiv operator:

Note that the regular version of the proxdiv operator, (2.17), is a special case of the stabilized variant with $\gamma=0$. With $K=\textnormal{get}K(\varepsilon)$ and $\mathcal{K}=\textnormal{get}\mathcal{K}(\hat{\alpha},\hat{\beta},\varepsilon)$ we observe that

For a threshold parameter $\tau>0$ we formally state the stabilized variant of Algorithm 1.

In the definitions for the stabilized kernel, (3.3b), and proxdiv-operator, (3.4), there still appear exponentials of the form $\exp(\cdot/\varepsilon)$, which may explode as $\varepsilon\to 0$. Extending the $\max$-argument trick in (3.1) to more general scaling algorithms entails similar questions. In the examples studied in Section 5 and those given in we find however, that evaluation of the exponential $\exp(-\gamma/\varepsilon)$ can be avoided. For the special case of standard optimal transport $\varepsilon$ no longer appears in the stabilized step.

2 ε𝜀\varepsilon-Scaling

It is empirically and theoretically well-known (cf. Section 1.1) that convergence of Algorithm 1 becomes slow as $\varepsilon\to 0$. A popular heuristic remedy is the so-called $\varepsilon$-scaling, where one subsequently solves the regularized problem with gradually decreasing values for $\varepsilon$. Let $\mathcal{E}=(\varepsilon_{1},\varepsilon_{2},\ldots,\varepsilon_{n})$ be a list of decreasing positive parameters. We extend Algorithm 2 as follows:

The dual variable $\beta$ is kept constant while changing $\varepsilon$, not the scaling factor $v$, because the optimal dual variables $(\alpha,\beta)$ usually have a stable limit as $\varepsilon\to 0$, while the scaling factors $(u,v)$ diverge (see Sect. 1.1 and also Theorem 4.16).

So far, very little is known theoretically about the behaviour of $\varepsilon$-scaling for the Sinkhorn algorithm (cf. Section 1.1). Empirically, it is shown in Sect. 5.2 that $\varepsilon$-scaling is highly efficient and the number of required iterations does not increase exponentially. We observe that indeed it behaves similar as in the auction algorithm, as discussed in . We work towards a theoretical quantification of this in Sect. 4.

Motivated by this, in practice we recommend a geometric decrease of $\varepsilon$ and choose $\varepsilon_{k}=\varepsilon_{0}\cdot\lambda^{k}$ such that $\varepsilon_{n}$ is the desired final value, $\varepsilon_{0}$ is on the order of the maximal values in the cost function $c$ and $\lambda\in(0,1)$ is a geometric scaling factor, typically in $[0.5,0.75]$. If $\lambda$ is too small, iterations will start far from convergence after each change of $\varepsilon$, increasing the risk of numerical instabilities and requiring more iterations. On the other hand, if $\lambda$ is too large, many stages of $\varepsilon$-scaling have to be performed, increasing numerical overhead.

3 Kernel Truncation

Storing the dense kernel $K$ and computing dense matrix multiplications during the scaling iterations (2.18) requires a lot of memory and time on large problems. For several problems with particular structure, remedies have been proposed (Sect. 1.1). But these do not comprise non-standard cost functions, as the one used for the Wasserstein-Fisher-Rao distance, (2.3). Moreover they are not compatible with the log-stabilization (Section 3.1), thus a certain level of blur cannot be avoided. We are looking for a more flexible method to accelerate solving.

For many unregularized transport problems the optimal coupling $\pi^{\dagger}$ is concentrated on a sparse subset of $X\times Y$. In fact, this is the underlying mechanism for the efficiency of most solvers discussed in Section 1.1. For the regularized problems the optimal coupling will usually be dense. This is due to the diverging derivative of the $\operatorname{KL}$ divergence at zero. However, as $\varepsilon\to 0$, the optimal coupling quickly converges to an unregularized solution (see Sect. 1.1, in particular [16, Thm. 5.8]). As $\varepsilon\to 0$, large parts of the coupling will approach zero exponentially fast.

So while we will not be able to exactly solve the full problem, by solving suitable sparse sub-problems, we may still expect a reasonable approximation. We formalize the concept of a sparse sub-problem.

Let $F_{X}$ and $F_{Y}$ be marginal functions and $c$ be a cost function as in Definition 2.1 and let $\mathcal{N}\subset X\times Y$. We introduce:

We call problems (2.4a) and (2.4b) with $c$ replaced by $\hat{c}$ the problems restricted to $\mathcal{N}$. This corresponds to adding the constraint $\operatorname{spt}\pi\subset\mathcal{N}$ to the primal problem, and only enforcing the constraint $\alpha(x)+\beta(y)\leq c(x,y)$ on $(x,y)\in\mathcal{N}$ in the dual problem. The entropy regularized variants of the restricted problems are obtained through replacing $K$ by $\hat{K}$ in (2.10a) and (2.10b).

Clearly, when $\mathcal{N}$ is sparse, then so is $\hat{K}$ and the restricted regularized problem can be solved faster and with less memory. We now quantify the error inflicted by restriction.

Let $\varepsilon>0$ and $\mathcal{N}\subset X\times Y$. Let $E$ and $J$ be unrestricted regularized primal and dual functionals with kernel $K$, as given in Definition 2.4, and let $\hat{E}$ and $\hat{J}$ be the functionals of the problems restricted to $\mathcal{N}$, with sparse kernel $\hat{K}$ (see Def. 3.1).

Further, let $(\alpha,\beta)$ be a pair of dual variables, let $u=\exp(\alpha/\varepsilon)$, $v=\exp(\beta/\varepsilon)$ be the corresponding scaling factors and let $\pi=\operatorname{diag}(u)\,\hat{K}\,\operatorname{diag}(v)$ be the corresponding (restricted) primal coupling.

Then we find for the primal-dual gap between $\pi$ and $(\alpha,\beta)$:

Together we obtain $E(\pi)-J(\alpha,\beta)=\hat{E}(\pi)-\hat{J}(\alpha,\beta)+\varepsilon\sum_{(x,y)\in(X\times Y)\setminus\mathcal{N}}u(x)\,K(x,y)\,v(y)$.

That is, the primal-dual gap for the original full functionals is equal to the gap for the truncated functionals plus the ‘mass’ that we have chopped off by truncating $K$ to $\hat{K}$, when using the scaling factors $u$ and $v$. If some $\mathcal{N}$ were known, on which most mass of the optimal $\pi^{\dagger}$ is concentrated, it would be sufficient to solve the problem restricted to $\mathcal{N}$, to get a good approximate solution. The remaining challenge is, how to identify $\mathcal{N}$ without knowing $\pi^{\dagger}$ before.

We propose an iterative re-estimation of $\mathcal{N}$, based on current dual iterates and to combine this with the log-stabilized iteration scheme (Section 3.1) and the computation of the stabilized kernel, (3.3b). For a threshold parameter $\theta>0$ we define the following functions:

$\textnormal{get}\hat{K}$ can be used instead of $\textnormal{get}\mathcal{K}$ in Algorithm 2. We refer to this as absorption iteration with truncation. For this combination one finds a simple bound for the primal-dual gap comparison of Proposition 3.2.

For $(x,y)\in(X\times Y)\setminus\mathcal{N}$ one has $\exp(-\tfrac{1}{\varepsilon}[c(x,y)-\hat{\alpha}(x)-\hat{\beta}(y)])<\theta$ and therefore

4 Multi-Scale Scheme

Finally, we propose to combine the stabilized sparse iterations with a hierarchical multi-scale scheme, analogous to the ideas in .

This serves two purposes: First, a hierarchical representation of the problem allows to determine the truncated sparse stabilized kernel $\textnormal{get}\hat{K}$, (3.9), with a coarse-to-fine tree search, without explicitly testing all pairs $(x,y)\in X\times Y$. The second reason is to make the combination of $\varepsilon$-scaling (Algorithm 3) with the truncated stabilized scheme more efficient. For a fixed threshold $\theta$, while $\varepsilon$ is large, the support of the truncated kernel $\textnormal{get}\hat{K}$ will contain many variables. At the same time, due to the blur induced by the regularization, the primal iterates will not provide a sharply resolved assignment. Solving the problems with large $\varepsilon$-value on a coarser grid reduces the number of required variables, without losing much spatial accuracy. As $\varepsilon$ decreases, so does the number of variables in $\textnormal{get}\hat{K}$ (since the exponential function decreases faster), and the resolution of $X$ and $Y$ can be increased. Therefore, it is reasonable to coordinate the reduction of $\varepsilon$ with increasing the spatial resolution of the transport problem, until the desired regularization and resolution are attained.

We will now briefly recall the hierarchical representation of a transport problem from .

For a discrete set $X$ a hierarchical partition is an ordered tuple $(\mathcal{X}_{0},\ldots,\mathcal{X}_{I})$ of partitions of $X$ where $\mathcal{X}_{0}=\{\{x\}\colon x\in X\}$ is the trivial partition of $X$ into singletons and each subsequent level is generated by merging cells from the previous level, i.e. for $i\in\{1,\ldots,I\}$ and any $\mathbf{x}\in\mathcal{X}_{i}$ there exists some $\hat{\mathcal{X}}\subset\mathcal{X}_{i-1}$ such that $\mathbf{x}=\bigcup_{\hat{\mathbf{x}}\in\hat{\mathcal{X}}}\hat{\mathbf{x}}$. For simplicity we assume that the coarsest level is the trivial partition into one set: $\mathcal{X}_{I}=\{X\}$. We call $I>0$ the depth of $\mathcal{X}$.

Then we define the extension $\hat{\alpha}=(\hat{\alpha}_{0},\ldots,\hat{\alpha}_{I})$ of $\alpha$ onto the full partition $\mathcal{X}$ by

for $i\in\{0,\ldots,I\}$ and $\mathbf{x}\in\mathcal{X}_{i}$ and analogous for $\hat{\beta}$ and $\beta$. Similarly, define an extension $\hat{c}$ of $c$ by

for $i\in\{0,\ldots,I\}$, $\mathbf{x}\in\mathcal{X}_{i}$ and $\mathbf{y}\in\mathcal{Y}_{i}$.

For $i\in\{0,\ldots,I\}$, $x\in\mathbf{x}\in\mathcal{X}_{i}$, $y\in\mathbf{y}\in\mathcal{Y}_{i}$ we find

Now we can implement a hierarchical tree-search for $\textnormal{get}\mathcal{N}$ (and analogously $\textnormal{get}\hat{K}$).

From (3.12) follows directly that Algorithm 4 implements (3.8).

The second purpose of the multi-scale scheme is the combination with $\varepsilon$-scaling. As explained above, the purpose is to reduce the number of variables while $\varepsilon$ is large. For an illustration see Fig. 1. For this, we divide the list $\mathcal{E}$ of regularization parameters $\varepsilon$ into multiple lists $(\mathcal{E}_{0},\ldots,\mathcal{E}_{I})$, with the largest values in $\mathcal{E}_{I}$ and the smallest (and final) values in $\mathcal{E}_{0}$, and sorted from largest to smallest within each $\mathcal{E}_{i}$. Then, for every $i$ from $I$ down to we perform $\varepsilon$-scaling with list $\mathcal{E}_{i}$ at hierarchical level $i$, using the dual solution at each level as initialization at the next stage. The full algorithm, combining $\log$-stabilization, $\varepsilon$-scaling, kernel truncation and the multi-scale scheme, is sketched next.

Note: ScalingAlgorithmStabilized refers to calling Algorithm 2 for solving the problem at scale $i$, with $\textnormal{get}\mathcal{K}$ replaced by $\textnormal{get}\hat{K}$, (3.9), with threshold $\theta$, implemented according to Algorithm 4. Accordingly, two arguments $i$ and $\theta$ were added. RefineDuals initializes the dual variables $(\alpha,\beta)$ at level $i$ by setting the values at $\mathbf{x}$ to the previous values at $\textnormal{parent}(\mathbf{x})$ for all cells $\mathbf{x}$ in $\mathcal{X}_{i}$.

To solve the problem at hierarchical scale $i$, not only do we need a coarse version of $c$, as given in (3.11). In addition we need hierarchical versions of the marginal functions $F_{X}$, $F_{Y}$, see (2.10). An appropriate choice is often clear from the context of the problem. For example, for an optimal transport problem between $\mu$ and $\nu$, see Def. 2.6, we set $F_{{\mathcal{X}}_{i}}=\iota_{\{\mu_{i}\}}$, where $\mu_{i}$ is taken from the multi-scale measure approximation of $\mu$ (see Def. 3.6). For the unbalanced transport problem with $\operatorname{KL}$ fidelity, Def. 2.3, we use $F_{\mathcal{X}_{i}}=\lambda\cdot\operatorname{KL}_{\mathcal{X}_{i}}(\cdot|\mu_{i})$.

This completes the modifications of the diagonal scaling algorithm. Their usefulness will be demonstrated numerically in Sect. 5.

Analogy between Sinkhorn and Auction Algorithm

In this section we develop a new complexity analysis of the Sinkhorn algorithm and examine the efficiency of $\varepsilon$-scaling. In an intuitive similarity between the Sinkhorn algorithm for the entropy regularized linear assignment problem and the auction algorithm was pointed out. This similarity motivates our approach.

In this section we only consider the standard Sinkhorn algorithm (as opposed to general scaling algorithms), since the auction algorithm solves the linear assignment problem and assumptions on fixed marginals $\mu$, $\nu$ are required for our analysis.

The auction algorithm is briefly recalled in Section 4.1. In Section 4.2 we introduce an asymmetric variant of the Sinkhorn algorithm, that is more similar to the original auction algorithm and provide an analogous worst-case estimate for the number of iterations until a given precision is achieved. A stability result for the dual optimal solutions under change of the regularization parameter $\varepsilon$ is given in Section 4.3 and we discuss how it relates to $\varepsilon$-scaling in Sect. 4.4.

For the sake of self-containedness, in this section we briefly recall the auction algorithm and its basic properties. Note that compared to the original presentation (e.g. ) we flipped the overall sign for compatibility with the notion of optimal transport.

The main loop of the auction algorithm is divided into two parts: During the bidding phase, elements of $X$ that are unassigned determine their locally most attractive counterpart in $Y$ (taking into account the current dual variables) and submit a bid for them. During the assignment phase, all elements of $Y$ that received at least one bid, pick the most attractive one and change the current assignment accordingly. A formal description is given in the following.

In the above algorithm, line 7 is usually replaced by $\alpha(x)\leftarrow\min_{y^{\prime}\in Y\setminus\{y\}}[c(x,y^{\prime})-\beta(y^{\prime})]$, which in practice may reduce the number of iterations. It does not affect the following worst-case analysis however, therefore we keep the simpler version.

We briefly summarize the main properties of the algorithm.

With $\varepsilon>0$ and $\beta^{(0)}=0_{Y}$, Algorithm 6 has the following properties:

$\alpha$ is increasing, $\beta$ is decreasing.

After each assignment phase one finds $\alpha(x)+\beta(y)\leq c(x,y)$ and $[\pi(x,y)>0]$ $\Rightarrow$ $[\alpha(x)+\beta(y)\geq c(x,y)-\varepsilon]$. The latter property is called $\varepsilon$-complementary slackness.

The primal iterate satisfies $\textnormal{P}_{X}\pi\leq\mu$ and $\textnormal{P}_{Y}\pi\leq\nu$.

The algorithm terminates after at most $N\cdot(C/\varepsilon+1)$ iterations, where $C=\max c$.

For a proof see for example . From $\varepsilon$-complementary slackness we deduce the following result.

Upon convergence, the primal-dual gap of $\pi$ and $(\alpha,\beta)$, cf. Def. 2.2, is bounded by $\left\langle c,\pi\right\rangle-(\left\langle\mu,\alpha\right\rangle+\left\langle\nu,\beta\right\rangle)\leq N\cdot\varepsilon$. If $c$ is integer and $\varepsilon<1/N$, then the final primal coupling is optimal.

During the auction algorithm it may happen that several elements in $X$ compete for the same target $y\in Y$, leading to the minimal decrease of $\beta(y)$ by $\varepsilon$ in each iteration. This phenomenon has been dubbed ‘price haggling’ and can cause poor practical performance of the algorithm, close to the worst-case iteration bound. The impact of price haggling can be reduced by the $\varepsilon$-scaling technique, where the algorithm is successively run with a sequence of decreasing values for $\varepsilon$, each time using the final value of $\beta$ as initialization of the next run (see also Algorithm 3). With this technique the factor $C/\varepsilon$ in the iteration bound can essentially be reduced to a factor $\log(C/\varepsilon)$. An analysis of the $\varepsilon$-scaling technique for more general min-cost-flow problems can be found in .

2 Asymmetric Sinkhorn Algorithm and Iteration Bound

We now introduce a slightly modified variant of the standard Sinkhorn algorithm, derive an iteration bound and make a comparison with the auction algorithm. We emphasize that this modification is primarily made to facilitate theoretical study of the algorithm and to understand why convergence becomes slow as $\varepsilon\to 0$. We do not advocate its merits in an actual implementation.

The only differences to the standard Sinkhorn algorithm (given by Algorithm 1 with proxdiv-operators (2.20)) lie in line 5 and in the choice of the specific stopping criterion (see Remark 4.13 for a discussion). In the standard algorithm one would set $v\leftarrow\hat{v}$. The modification implies that $v$ is monotonously decreasing, which implies the following result for Algorithm 7 in the spirit of Proposition 4.2. This monotonicity is crucial for bounding the number of iterations (see also Remark 4.15).

$u$ and $\alpha=\varepsilon\,\log u$ are increasing, $v$ and $\beta=\varepsilon\,\log v$ are decreasing, $q$ is increasing.

$\textnormal{P}_{X}\,\pi\leq\mu$ and $\textnormal{P}_{Y}\,\pi\leq\nu$. We say $\pi$ is sub-feasible.

There exists some $y^{\ast}\in Y$ such that $v(y^{\ast})=v^{(0)}(y^{\ast})$ for all iterations.

With these tools we can bound the total number of iterations to reach a given precision.

Let us look at the first ‘bid’ $\alpha^{(1)}$. With $c\geq 0$ we have

Combining this with (4.2) we obtain $n\leq 1+\tfrac{C}{\varepsilon\,(1-q^{(n)})}$. So, as long as $q^{(n)}<q_{\textnormal{target}}$ we have $n<1+\tfrac{C}{\varepsilon\,(1-q_{\textnormal{target}})}$. By contraposition we know that there is some $n\leq 2+\tfrac{C}{\varepsilon\,(1-q_{\textnormal{target}})}$ such that $q^{(n)}\geq q_{\textnormal{target}}$.

And finally, we formally establish convergence of the iterates.

The criterion $q\geq q_{\textnormal{target}}$ is motivated by Lemma 4.7, to provide a minimal increment of $\alpha$ during iterations. $1-q$ measures the mass that is still missing and is equal to the $L^{1}$ error between the marginals of $\pi$ and the desired marginals $\mu$ and $\nu$. In pathological cases the dual variables $(\alpha,\beta)$ may still be far from optimizers, even though $q\geq q_{\textnormal{target}}$ (see Example 4.14). In [21, Lemma 2] linear convergence of the marginals in the Hilbert projective metric is proven. This is a stricter measure of convergence, less prone to ‘premature’ termination. However, for small $\varepsilon$ the contraction factor is roughly $1-4\,\exp(-C/\varepsilon)$, which is impractical. The scaling $\mathcal{O}(1/\varepsilon)$ predicted by Proposition 4.9 is consistent with numerical observations when one uses the $L^{1}$ or $L^{\infty}$ marginal error as stopping criterion (Sect. 5.2). Therefore we consider the $q$-criterion to be a reasonable measure for convergence, as long as one keeps $1-q_{\textnormal{target}}\ll\delta$ (Example 4.14).

We consider the $1\times 2$ toy problem with the following parameters:

for some $C>0$, $\delta\in(0,1)$ and some regularization strength $\varepsilon>0$. And we consider the scaling factors (one for $X$, two choices for $Y$): $u=\begin{pmatrix}1\end{pmatrix}^{\top}$, $v_{1}=\begin{pmatrix}1&1\end{pmatrix}^{\top}$, $v_{2}=\begin{pmatrix}1&e^{C/\varepsilon}\end{pmatrix}^{\top}$. Let $\pi_{i}=\operatorname{diag}(u)\,K\,\operatorname{diag}(v_{i})$ and corresponding total masses $q_{i}$, $i=1,2$. We find:

$\pi_{2}$ and $(\alpha,\beta_{2})=\varepsilon\,\log(u,v_{2})$ are primal and dual solutions. $\pi_{1}$ is sub-feasible (see Proposition 4.5). For fixed $\varepsilon>0$, as $\delta\to 0$, $q_{1}$ tends to 1 (but is strictly smaller), i.e. the pair $(u,v_{1})$ has almost converged in the $q$-measure sense, but the distance between $\beta_{1}=\varepsilon\,\log v_{1}$ and the actual solution $\beta_{2}$ is $C$.

For now assume $|X|=|Y|=N$ and $\mu$, $\nu$ are normalized counting measures. Then line 4 in Algorithm 7, expressed in dual variables, becomes

These are formally similar to the corresponding lines 7 and 13 in Algorithm 6. We can interpret the $u$-update in Algorithm 7 as $x$ not just submitting a bid to the best candidate $y$, but to all candidates, weighted by the attractiveness (recall that in the Sinkhorn algorithm, a change in the dual variable directly implies a change in the primal iterate via (2.11)). Conversely, in line 5, $y$ does not only accept the best bid, but bids from all candidates, again weighted by price. If there are too many bids (i.e. if $\hat{v}(y)<v(y)$), $\beta(y)$ decreases and thereby rejects superfluous offers.

Consequently, in Algorithm 7 one can observe that points in $X$ compete for the mass in $Y$ in a way similar to the auction algorithm by repeatedly increasing their prices until a different target seems more attractive or other competitors lose interest. In both algorithms the minimal increment is related to the parameter $\varepsilon$ which leads to iteration bounds that are proportional to $1/\varepsilon$ (Props. 4.2 and 4.9). An attempt to mimic the analysis of $\varepsilon$-scaling is made in Section 4.4 (cf. Remark 4.26).

One can interpret the standard Sinkhorn algorithm with $v\leftarrow\hat{v}$ as $y$ submitting a ‘counter-bid’ if it has not received enough bids. Such a symmetrization has also been discussed for the auction algorithm. But then the complexity analysis based on monotonous dual variables breaks down, and the algorithm may even run indefinitely (see ‘down iterations’ in ).

3 Stability of Dual Solutions

The main result of this section is Theorem 4.16, which provides stability of dual solutions to entropy regularized optimal transport (Def. 2.6) under changes of the regularization parameter $\varepsilon$. Its implications for $\varepsilon$-scaling are discussed in Sect. 4.4.

studies the convergence of entropy regularized linear programs to the unregularized variant and can be used to understand the limit of entropy regularized optimal transport (Def. 2.6). To apply , the constraint matrix must have full rank and the set of optimal solutions to (2.5b) must be bounded. When the cost $c$ is finite this is achieved by arbitrarily fixing one dual variable, e.g. $\alpha(x_{0})=0$, and removing the corresponding column from the dual constraint matrix. The slight difference in the definition of the entropy (or the dual exponential barrier) can be absorbed into a change of variables which converges to the identity in the limit $\varepsilon\to 0$.

Then [16, Props. 3.1 and 3.2] imply that the optimal solutions of (2.19b) remain bounded and converge to a particular solution of (2.5b) as $\varepsilon\to 0$. Furthermore, provides statements about the convergence of the optimal couplings (Prop. 4.1) and the asymptotic behaviour (Thm. 5.8).

The bounds derived in depend on the geometry of the primal and dual feasible polytopes of (2.5), i.e. on the transport cost function $c$. In contrast, the bound of Thm. 4.16 does not depend on $c$. The motivation for deriving such a bound is the implication for $\varepsilon$-scaling, see Section 4.4. Note that Thm. 4.16 also implies that the optimal dual variables remain bounded as $\varepsilon\to 0$.

The proof requires several auxiliary definitions and lemmas. The estimate consists of two contributions: One stems from following paths within connected components of what we call assignment graph (defined in the following lemma), using the primal-dual relation (2.11). This reasoning is analogous to the proof strategy for $\varepsilon$-scaling in the auction algorithm (see ). However, between different connected components (2.11) is too weak to yield useful estimates. So a second contribution arises from a stability analysis of effective diagonal problems (in Lemmas 4.21 and 4.23).

For two feasible couplings $\pi_{1}$, $\pi_{2}\in\Pi(\mu,\nu)$ and a threshold $M^{-1}\geq 1$ the corresponding assignment graph is a bipartite directed graph with vertex sets $(X,Y)$ and the set of directed edges

where $(a,b)\in\mathcal{E}$ indicates a directed edge from $a\to b$.

The assignment graph has the following properties:

Every node has at least one incoming and one outgoing edge.

Let $X_{0}\subset X$, $Y_{0}\subset Y$ such that there are no outgoing edges from $(X_{0},Y_{0})$ to the rest of the vertices, then $|\mu(X_{0})-\nu(Y_{0})|<1/M$. This is also true when there are no incoming edges from the rest of the vertices. If $\mu$ and $\nu$ are atomic, with atom size $1/M$ (see Assumption 1), then $\mu(X_{0})=\nu(Y_{0})$.

Assume, a node $x\in X$ had no outgoing edge. Then $\sum_{y\in Y}\pi_{2}(x,y)<\mu(x)/M\leq\mu(x)$. This contradicts $\pi_{2}\in\Pi(\mu,\nu)$. Existence of incoming edges follows analogously.

Let $\hat{X}_{0}=X\setminus X_{0}$, $\hat{Y}_{0}=Y\setminus Y_{0}$. If $(X_{0},Y_{0})$ has no outgoing edges, then

Since $\pi_{1}$, $\pi_{2}\in\Pi(\mu,\nu)$, the first inequality implies $\mu(X_{0})=\pi_{2}(X_{0}\times Y)=\pi_{2}(X_{0}\times Y_{0})+\pi_{2}(X_{0}\times\hat{Y}_{0})<\nu(Y_{0})+1/M$ and the second inequality implies $\nu(Y_{0})<\mu(X_{0})+1/M$, i.e. $|\mu(X_{0})-\nu(Y_{0})|<1/M$. With Assumption 1 for atom size $1/M$, this implies $\mu(X_{0})=\nu(Y_{0})$. The statement about incoming edges follows from $\mu(\hat{X}_{0})=1-\mu(X_{0})$ and $\nu(\hat{Y}_{0})=1-\nu(Y_{0})$.

Every node in $(X,Y)$ is part of at least one strongly connected component (containing at least the node itself). If two strongly connected components have a common element, they are identical. Hence, the strongly connected components form partitions of $X$ and $Y$. For some $x\in X$ (or $y\in Y$), let $X_{\textnormal{out}}\subset X$ and $Y_{\textnormal{out}}\subset Y$ be the set of nodes that can be reached from $x$, let $X_{\textnormal{in}}\subset X$ and $Y_{\textnormal{in}}\subset Y$ be the set of nodes from which one can reach $x$ and let $(X_{\textnormal{con}}=X_{\textnormal{out}}\cap X_{\textnormal{in}},Y_{\textnormal{con}}=Y_{\textnormal{out}}\cap Y_{\textnormal{in}})$ be the strongly connected component of $x$. Clearly $(X_{\textnormal{out}},Y_{\textnormal{out}})$ has no outgoing edges. Hence, by (ii) one has $\mu(X_{\textnormal{out}})=\nu(Y_{\textnormal{out}})$. Moreover, $(X_{\textnormal{out}}\setminus X_{\textnormal{in}},Y_{\textnormal{out}}\setminus Y_{\textnormal{in}})$ has no outgoing edges, hence $\mu(X_{\textnormal{out}}\setminus X_{\textnormal{in}})=\nu(Y_{\textnormal{out}}\setminus Y_{\textnormal{in}})$, from which follows that $\mu(X_{\textnormal{con}})=\nu(Y_{\textnormal{con}})$.

where $J$ denotes the dual functional of entropy regularized optimal transport (2.19b), and

Since maximizing $\hat{J}$ corresponds to maximizing $J$ over an affine subspace, clearly $\hat{\beta}^{\dagger}$ is a maximizer of $\hat{J}$. Since $\hat{J}$ inherits the invariance of $J$ under constant shifts, any $\hat{\beta}^{\dagger\dagger}$ of the form given above, is also a maximizer. Consequently, we may add the constraint $\hat{\beta}(1)=0$, which does not change the optimal value. With this added constraint the functional becomes strictly convex, which implies a unique optimizer. Hence, any optimizer of the unconstrained functional can be written in the form of $\hat{\beta}^{\dagger\dagger}$.

Let us now give a more explicit expression of $\hat{J}(\hat{\beta})$. We find

Minimizers of $\hat{J}_{\varepsilon,d}$ exist.

That is, $\operatorname{maxdiam}(w)$ is the length of the longest cycle-less path in $\{1,\ldots,R\}$ with edge lengths $w$.

The proofs of Theorem 4.16 and Lemma 4.23 can be found in Appendix A.

4 Application To ε𝜀\varepsilon-Scaling

Assuming that we know the dual solution for some $\varepsilon_{1}>0$, then Theorem 4.16 allows to bound the number of iterations of Algorithm 7 for some smaller $\varepsilon_{2}\in(0,\varepsilon_{1})$, independently of bounds on the cost function $c$. This may have implications for the efficiency of $\varepsilon$-scaling (see Remark 4.26).

Consider the set-up of Theorem 4.16. In particular, let $\varepsilon_{1}>\varepsilon_{2}>0$ be two regularization parameters, let $(\alpha_{1},\beta_{1})$, $(\alpha_{2},\beta_{2})$ be corresponding optimizers of (2.19b). If Algorithm 7 is initialized with $v^{(0)}=\exp(\beta_{1}/\varepsilon_{2})$, with regularization $\varepsilon_{2}$, and for a given $q_{\textnormal{target}}\in(0,1)$, the number of iterations $n$ necessary to achieve $q^{(n)}\geq q_{\textnormal{target}}$ is bounded by

For the optimal scaling factor $u_{1}$ of the $\varepsilon_{1}$-problem we find:

This implies $u_{1}(x)^{-1}\,\nu(y)^{-1}\geq\exp\left(-\tfrac{1}{\varepsilon_{1}}[c(x,y)-\beta_{1}(y)]\right)$ for all $(x,y)\in X\times Y$. With this we can bound the first iterate of the $\varepsilon_{2}$-run of the algorithm by:

where we have used $\nu(y)\geq 1/M$, Assumption 1. Eventually we find $\alpha^{(1)}(x)\geq\alpha_{1}(x)-\varepsilon_{1}\,\log M$.

Now we combine Algorithm 7 with $\varepsilon$-scaling, (cf. Algorithm 3). For $\varepsilon=\hat{\varepsilon}\cdot\lambda^{-k}\geq C$, according to Proposition 4.9 it will take at most $2+\frac{1}{1-q_{\textnormal{target}}}$ iterations. It is tempting to deduce from Proposition 4.24 that for each subsequent value of $\varepsilon$ at most $2+\frac{A}{\lambda\,(1-q_{\textnormal{target}})}$ iterations are required, with $A=N\cdot(4\log N+24\log M)+\log M$. For $N>1$ the total number of iterations would then be bounded by $(2+\frac{A}{\lambda\,(1-q_{\textnormal{target}})})\cdot(k+1)$. For fixed $\lambda$ the step parameter $k$ scales like $\log(C/\hat{\varepsilon})$. Consequently, the total number of iterations would be bounded by $\mathcal{O}(\log(C/\hat{\varepsilon}))$ w.r.t. the cost function and regularization, which would be analogous to $\varepsilon$-scaling for the auction algorithm (Remark 4.4).

There is an obvious gap in this reasoning: Theorem 4.16 assumes that $\beta_{1}$ is known exactly, while Algorithm 7 only provides an approximate result. From Example 4.14 we learn that in extreme cases this difference can be substantial and disrupt the efficiency of $\varepsilon$-scaling. Thus, additional assumptions on the problem are required to make the above argument rigorous.

However, as discussed in Remark 4.13, in practice we usually observe that approximate iterates are sufficient and we can therefore hope that $\varepsilon$-scaling does indeed serve its purpose.

Numerical Examples

Now we present a series of numerical experiments to confirm the usefulness of the modifications proposed in Sect. 3. We show that runtime and memory usage are reduced substantially. At the same time the adapted algorithm is still as versatile as the basic version of , Algorithm 1. But Algorithm 5 can solve larger problems at lower regularization, yielding very sharp results. We give examples for unbalanced transport, barycenters and Wasserstein gradient flows. The code used for the numerical experiments is available from the author’s website.https://github.com/bernhard-schmitzer

2 Efficiency of Enhanced Algorithm

The numerical efficiency of the subsequent modifications presented in Sect. 3, applied to the standard Sinkhorn algorithm, is illustrated in Fig. 2. While the stabilized algorithm (i) is not yet faster than the naive implementation, it can robustly solve the problem for all given values of $\varepsilon$. The required number of iterations scales like $\mathcal{O}(1/\varepsilon)$, in good agreement with the complexity analysis of Sect. 4.2. With $\varepsilon$-scaling (ii) the number of iterations is decreased substantially. Replacing the dense kernel with the adaptive truncated sparse kernel (iii) does not change the number of required iterations, but saves time and memory. With the multi-scale scheme the required number of iterations is slightly increased, since the initial dual variables obtained at a coarser level are only approximate solutions. But by reducing the number of variables during the early $\varepsilon$-scaling stages, the runtime is further decreased (cf. Fig. 1). The combination of all modifications leads to an average total speed-up of more than two orders of magnitude on this problem type.

A runtime benchmark and study of the sparsity of the truncated kernel are given in Fig. 3. The runtime scales approximately linear with $|X|$ and for large problems the algorithm becomes faster than the adaptive sparse linear programming solver . The final number of variables in the sparse kernel is comparable with the number of variables in , for higher values of $\varepsilon$, during scaling, more memory is required (cf. Fig. 4). This underlines again the importance of the coarse-to-fine scheme (Sect. 3.4). It should be noted, that Fig. 2 shows results for $64\times 64$ images, the smallest image size in Fig. 3. For larger images the runtime difference between (i-iv) would be even larger, but due to time and memory constraints, only (iv) can be run practically.

The impact of different final values for $\varepsilon$ is outlined in Fig. 4. As expected, the number of variables in the truncated kernel increases with $\varepsilon$. This leads to two competing trends in the overall runtime: For large $\varepsilon$, the kernel truncation is less efficient, leading to an increase with $\varepsilon$. For small $\varepsilon$, the number of variables is very small, but more and more stages of $\varepsilon$-scaling are necessary, increasing the runtime as $\varepsilon$ decreases further. Convergence of the regularized optimal dual variables to the unregularized optimal duals is exemplified in the right panel, justifying the use of the approximate entropy regularization technique for transport-type problems. While one may consider the dual sub-optimality at $\varepsilon\approx 30\,h^{2}$ sufficiently accurate, we point out that the corresponding primal coupling still contains considerable blur (cf. Fig. 1) and that due to less sparsity the runtime is actually higher than for $\varepsilon\approx h^{2}$.

As illustrated by Figs. 3 and 4, by choosing the threshold for the stopping criterion and the desired final $\varepsilon$, one can tune between required precision and available runtime.

The numerical findings presented in Figs. 2-4 underline how each of the modifications discussed in Sect. 3 builds on the previous ones and that all four of them are required for an efficient algorithm. The log-domain stabilization is an indispensable prerequisite for running the scaling algorithms with small regularization. However, for small $\varepsilon$, convergence tends to become extremely slow (cf. Fig. 2), therefore $\varepsilon$-scaling is needed to reduce the number of iterations. For small $\varepsilon$, kernel truncation significantly reduces the number of variables and accelerates the algorithm (cf. Figs. 2 and 4). However, for large $\varepsilon$ (which must be passed during $\varepsilon$-scaling), far fewer variables are truncated and the algorithm cannot be run on large problems. This can be avoided by using the coarse-to-fine scheme, completing the algorithm. In principle it is possible, only to combine log-domain stabilization with kernel truncation, and to skip $\varepsilon$-scaling and the coarse-to-fine scheme. While this tends to solve the stability and memory issues, convergence is still impractically slow.

3 Versatility

The framework of scaling algorithms developed in , see Sect. 2, allows to solve more general transport-type problems for which the enhancements of Sect. 3 still apply. We now give some examples to demonstrate this flexibility. The scope of the following examples is similar to , but with Algorithm 5 one can solve larger problems with smaller regularization.

A proof is given in . Compared to the standard Sinkhorn algorithm, the only modification is the pointwise power of the iterates. As $\lambda\to\infty$ the Sinkhorn iterations are recovered. In the stabilized operator only the exponential \exp\big{(}\tfrac{-\alpha}{\lambda+\varepsilon}\big{)} needs to be evaluated, which remains bounded as $\varepsilon\to 0$. Algorithm 5 performs similarly with $\operatorname{KL}$-fidelity as with fixed marginal constraints, allowing to efficiently solve large unbalanced transport problems. Since the truncation scheme can also be used with non-standard cost functions such as (2.3), this includes in particular the Wasserstein-Fisher-Rao (WFR) distance. Fig. 5 shows a geodesic for the WFR distance, to intuitively illustrate its properties. The geodesic has been computed as weighted barycenters between its endpoints (see below). For a direct dynamic formulation we refer to . For the relation to the $\operatorname{KL}$ soft-marginal formulation, Def. 2.3, see .

Wasserstein barycenters

Wasserstein barycenters as a natural generalization of the Riemannian center of mass have been studied in . The computation of entropy regularized Wasserstein barycenters with a Sinkhorn-type scaling algorithm has been presented in , an alternative numerical approach can be found in . The iterations can be considered as a special case of the framework in . Here, we very briefly recall the iterations. Derivations and proofs can be found in .

and $K$ is the kernel (2.8) over $X\times X$ for the cost $c=d^{2}$. When an optimizer $(\pi^{\dagger}_{i})_{i}$ is found, the common second marginal of all $\pi^{\dagger}_{i}$ is the sought-after barycenter. To solve (5.2) one considers again a suitable dual problem and uses alternating optimization. Updates corresponding to $F_{1}$ decompose into independent standard Sinkhorn iterations for each marginal, the update for $F_{2}$ couples all marginals, see . The adaptations from Sect. 3 remain applicable. A barycentric triangle computed with Algorithm 5 is shown in Fig. 6. The log-domain stabilization allows to reach a lower final regularization $\varepsilon$ as for example in . Regularization can be made so small that discretization artifacts become visible. While this may not look entirely pleasing, it clearly gives a better approximation to the unregularized problem and illustrates that with log-domain stabilization entropy regularized numerical methods can produce sharp results.

Wasserstein-Fisher-Rao barycenters

Similarly one can define barycenters for transport distances with $\operatorname{KL}$ marginal fidelity, which includes the Gaussian Hellinger-Kantorovich (GHK) distance and the Wasserstein-Fisher-Rao (WFR) distance (Def. 2.3). The primal functional is given by (5.2) with

where $\Lambda>0$ is a global weight of the $\operatorname{KL}$-fidelity. When a primal optimizer is found, the minimizing $\sigma$ in $F_{2}$ yields the sought-after barycenter. We refer to for details. Partial optimization corresponding to $F_{1}$ can again be done separately for each marginal, leading to $\operatorname{KL}$ fidelity updates as given by (5.1). The update corresponding to $F_{2}$ is again coupled , adaptations from Sect. 3 remain applicable.

Wasserstein Gradient Flows

In diagonal scaling algorithms were extended to compute proximal steps for entropy regularized optimal transport to approximate gradient flows in Wasserstein space (cf. Sect. 1.1). This was then subsumed into the general framework of . Here we given an example for the porous medium equation, for more details we refer to . Let

Conclusion

Scaling algorithms for entropy regularized transport-type problems have become a wide-spread numerical tool. Naive implementations have some severe numerical limitations, in particular for small regularization and on large problems. In this article, we proposed an enhanced variant of the standard scaling algorithm to address these issues: Diverging scaling factors and slow convergence are remedied by log-domain stabilization and $\varepsilon$-scaling. Required runtime and memory are significantly reduced by adaptive kernel truncation and a coarse-to-fine scheme. A new convergence analysis for the Sinkhorn algorithm was developed. Numerical examples showed the efficiency of the enhanced algorithm, confirmed the scaling predicted by the convergence analysis and demonstrated that the algorithm can produce sharp results on a wide range of transport-type problems. Potential directions for future research are the more detailed study of $\varepsilon$-scaling, a more systematic understanding of the stability of the log-domain stabilization and application to multi-marginal problems.

Lénaïc Chizat, Luca Nenna and Gabriel Peyré are thanked for stimulating discussions. Bernhard Schmitzer was supported by the European Research Council (project SIGMA-Vision).

Appendix A Additional Proofs

The first order optimality condition for the functional yields for the $i$-th component of $\beta$:

From the optimality conditions for $\beta_{a}(i_{1})$, $a=1,2$, and (1.3) we obtain:

where $w(i,j)=\max\{-\Delta d(i,j),\Delta d(j,i)\}$. This implies there is some $i_{2}\in\{1,\ldots,R\}\setminus\{i_{1}\}$ with

where $W(i_{1},k)=\max\{-\Delta D(i_{1},k),\Delta D(k,i_{1})\}$ for $k\in I_{2}$ and $\Delta D=D_{2}-D_{1}$. With (1.3) we find

and eventually $W(i_{1},k)\leq\max_{j\in I_{1}}(w(j,k)-\Delta\beta(j))+\Delta\beta(i_{1})+\max\{\varepsilon_{1},\varepsilon_{2}\}\cdot\log R$. So there is some index $i_{3}\in I_{2}$ such that

The index $i_{3}$ will be called a child of the minimizing index $j\in I_{1}$ on the r.h.s. (or one of the minimizing indices). Then we add $i_{3}$ to the set $I_{1}$ and repeat the argument with the reduced functional, to obtain an index $i_{4}$ and repeat this until $I_{1}$ contains all indices.

Since we assign every new index $i_{k}$ that is added to $I_{1}$ as a child to one parent node in $I_{1}$, this also constructs a tree graph with root node $i_{1}$ (finiteness of $d$ and consequently $D$ implies that this graph is connected). For an index $i_{k}$ let $(i_{1},i_{2},\ldots,i_{k})$ be the unique path from the root to $i_{k}$. Then

Since $\Delta\beta(i_{1})=\max\Delta\beta$ the result follows.

A.2 Proof of Theorem 4.16

Let $\pi_{1}$, $\pi_{2}$ be the primal optimizers associated with $(\alpha_{1},\beta_{1})$ and $(\alpha_{2},\beta_{2})$ and consider the assignment graph for $\pi_{1}$ and $\pi_{2}$ and threshold $1/M$ (see Lemma 4.19). Let $\{(X_{i},Y_{i})\}_{i=1}^{R}$ be the strongly connected components of the assignment graph. By virtue of Lemma 4.19(iii), $\mu(X_{i})=\nu(Y_{i})$ for $i=1,\ldots,R$. Pick some representatives $\{y_{i}\}_{i=1}^{R}\subset Y$ such that $y_{i}\in Y_{i}$ for $i=1,\ldots,R$.

Now we derive some estimates on $\Delta d$. Consider once more the assignment graph for $\pi_{1}$, $\pi_{2}$ and threshold $1/M$. For every edge $y\to x$ we have (using (2.11))

Moreover, from the marginal conditions we find $\pi_{2}(x,y)\leq\nu(y)$, which implies

Combining the two estimates, we obtain $\Delta\alpha(x)+\Delta\beta(y)\leq(\varepsilon_{1}+\varepsilon_{2})\,\log M\leq 2\varepsilon_{1}\,\log M:=L$. Similarly, for edges $x\to y$ we obtain $\Delta\alpha(x)+\Delta\beta(y)\geq-(\varepsilon_{1}+\varepsilon_{2})\,\log M\geq-L$. Let now $(y_{1},x_{1},\ldots,y_{k})$ be an alternating path in $(X,Y)$, then, by combining the above inequalities we find $\Delta\beta(y_{j+1})\geq\Delta\beta(y_{j})-2\cdot L$ for $j=1,\ldots,k-1$ and eventually

Similarly, for a path $(x_{1},y_{2},x_{2},\ldots,y_{k})$ get $\Delta\alpha(x_{1})+\Delta\beta(y_{k})\geq-(2\,k-1)\cdot L$, and for a path $(y_{1},x_{1},\ldots,y_{k},\allowbreak x_{k})$ get $\Delta\alpha(x_{k})+\Delta\beta(y_{1})\leq(2\,k-1)\cdot L$.

where we used $|\Delta\varepsilon\log(\mu(x)\,\nu(y))|\leq 2\varepsilon_{1}\,\log M=L$. From this follows that $w(i,j)\leq 8\,\max\{|Y_{i}|,|Y_{j}|\}\cdot\varepsilon_{1}\,\log M+2\,\varepsilon_{1}\,\log N$, which in turn implies that $\operatorname{maxdiam}(w)\leq 16\,\varepsilon_{1}\,N\,\log M+2\,\varepsilon_{1}\,R\,\log N$.

Recall that $\Delta\hat{\beta}=\hat{\beta}_{2}^{\dagger}-\hat{\beta}_{1}^{\dagger}$, where $\hat{\beta}_{a}^{\dagger}$, $a=1,2$, are the optimizers of the effective diagonal problems. Then from Lemma 4.23, and by bounding $R\leq N$ we obtain that

and analogously we get the equivalent bound for $\Delta\alpha$.

References