Randomized methods for computing optimal transport without regularization and their convergence analysis

The optimal transport problem was first introduced by Monge in 1781, which aims to find the most cost-efficient way to transport mass from a set of sources to a set of sinks. Later, the theory was modernized and revolutionized by Kantorovich in 1942, who found a key link between optimal transport and linear programming. In recent years, optimal transport has become a popular and powerful tool in data science, especially in image processing, machine learning, and deep learning areas, where it provides a very natural way to compare and interpolate probability distributions. For instance, in generative models arjovsky2017wasserstein ; lei2019geometric ; wang2022deepparticle , a natural penalty function is the Wasserstein distance (a concept closely related to OT) between the data and the generated distribution. In image processing, the optimal transport plan, which minimizes the transportation cost, provides solutions to image registration haker2004optimal and seamless copy perrot2016mapping . Apart from data science, in the past three decades, there has been an explosion of research interest in optimal transport because of the deep connections between the optimal transport problems with quadratic cost functions and a diverse class of partial differential equations (PDEs) arising in statistical mechanics and fluid mechanics; see e.g. brenier1991polar ; benamou2000computational ; otto2001geometry ; jordan1998variational ; villani2021topics for just a few of the most prominent results and references therein.

Inspired by this research progress, we have developed efficient numerical methods to solve multi-scale PDE problems using the optimal transport approach. Specifically, in our recent paper, we proposed a deep particle method for learning and computing invariant measures of parameterized stochastic dynamical systems wang2022deepparticle . To achieve this goal, we designed a deep neural network (DNN) to map a uniform distribution (source) to an invariant measure (target), where the Péclet number is an input parameter for the DNN. The network is trained by minimizing the 2-Wasserstein distance ($W_{2}$) between the measure of network output $\mu$ and target measure $\nu$. We consider a discrete version of $W_{2}$ for finitely many samples of $\mu$ and $\nu$, which involves a linear program (LP) optimized over doubly stochastic matrices sinkhorn1964relationship .

Solving the LP directly using the interior point method wright1997primal is too costly. Motivated by the domain decomposition method toselli2004domain in scientific computing, which solves partial differential equations using subroutines that solve problems on subdomains and has the advantage of saving memory (i.e., using the same computational resource, it can compute a larger problem), we devised a mini-batch interior point method. This approach involves sampling smaller sub-matrices while preserving row and column sums. It has proven to be highly efficient and integrates seamlessly with the stochastic gradient descent (SGD) method for overall network training. However, we did not obtain convergence analysis for this mini-batch interior point method in our previous work wang2022deepparticle .

The objectives of this paper are twofold. First, we aim to provide rigorous convergence analysis for the mini-batch interior point method presented in wang2022deepparticle , with minimal modifications. Second, we seek to enhance the mini-batch selection strategy, thereby achieving improved and more robust performance in computing optimal transport problems. We recognize that the mini-batch interior point method aligns with the random block coordinate descent (RBCD) method in optimization terminology. Specifically, it applies the block coordinate descent (BCD) method to the LP problem directly, selects the working set randomly, and solves subproblems using the primal-dual interior point method wright1997primal or any other efficient linear programming solver. Encouraged by the demonstrated efficiency of this approach, we will develop theoretical results for solving LP with RBCD methods and explore various strategies for selecting working sets.

In this work, we first introduce an expected Gauss-Southwell-$q$ rule to guide the selection of the working set. It enables almost sure convergence and a linear convergence rate in expectation when solving a general standard LP. Based on this rule, we develop a vanilla RBCD method - RBCD₀, which selects the working set with complete randomness. Then, we investigate the special linear system present in the LP formulation of OT. We characterize all the elementary vectors of the null space and provide a strategy for finding the conformal realization of any given vector in the null space at a low computational cost. Based on these findings, we propose various approaches to refine the working set selection and improve the performance of RBCD₀. A better estimation of the constant in the linear convergence rate is shown. Moreover, we incorporate an acceleration technique inspired by the momentum concept to improve the algorithm’s efficiency. Inexact versions of the RBCD methods are also discussed.

We perform numerical experiments to evaluate the performance of the proposed methods. Synthetic data sets of various shapes/dimensions and invariant measures generated from IPM methods are utilized to create distributions. Our experiments first compare different RBCD methods proposed in this paper, demonstrating the benefits of refining working set selection and verifying the effectiveness of the acceleration technique. We also illustrate the gap between theory and practice regarding convergence rate, sparse solutions generated by the proposed RBCD methods, and discuss the choice of subproblem size. In the second set of experiments, we compare the best-performance method, ARBCD, with Sinkhorn’s algorithm. Preliminary numerical results show that ARBCD is comparable to Sinkhorn’s algorithm in computation time when seeking solutions with relatively high accuracy. ARBCD is also comparable to a recently proposed interior point inspired algorithm in memory-saving settings. We also test ARBCD on a large-scale OT problem, where Gurobi runs out of memory. This further justifies the memory-saving advantage of ARBCD.

BCD and RBCD are well-studied for essentially unconstrained smooth optimization (sometimes allow separable constraints or nonsmooth separable objective functions): beck2013convergence ; gurbuzbalaban2017cyclic ; sun2021worst investigate BCD with cyclic coordinate search; nesterov2012efficiency ; lu2015complexity ; richtarik2014iteration study RBCD to address problems with possibly nonsmooth separable objective functions; other related works include theoretical speedup of RBCD (richtarik2016parallel ; necoara2016parallel ), second-order sketching (qu2016sdna ; berahas2020investigation ). However, much less is known for their convergence properties when applied to problems with nonseparable nonsmooth functions as summands or coupled constraints. To our best knowledge, no one has ever considered using the RBCD to solve general LP before and the related theoretical guarantees are absent. In necoara2017random , the authors studied the RBCD method to tackle problems with a convex smooth objective and coupled linear equality constraints $x_{1}+x_{2}+\ldots+x_{N}=0$; a similar algorithm named random sketch descent method necoara2021randomized is investigated to solve problems with a general smooth objective and general coupled linear equality constraints $Ax=b$. However, after adding the simple bound constraints $x\geq 0$, the analysis in necoara2017random ; necoara2021randomized may not work anymore, nor can it be easily generalized. Beck beck20142 studied a greedy coordinate descent method but focused on a single linear equality constraint and bound constraints. In Paul Tseng and his collaborators’ work tseng2009block ; tseng2009coordinate ; tseng2010coordinate , a block coordinate gradient descent method is proposed to solve linearly constrained optimization problems including general LP. In these works, a Gauss-Southwell-$q$ rule is proposed to guide the selection of the working set in each iteration. Therefore, the working set selected in a deterministic fashion can only be decided after solving a quadratic program with a similar problem size as the original one. In contrast, our proposed mini-batch interior point/RBCD method approach selects the working set through a combination of randomness and low computational cost. Another research direction that addresses separable functions, linearly coupled constraints, and additional separable constraints involves using the alternating direction method of multipliers (ADMM) chen2016direct ; he20121 ; xie2019si ; xie2021tractable . This method updates blocks of primal variables in a Gauss-Seidal fashion and incorporates multiplier updates as well.

Encouraged by the success in applying Sinkhorn’s algorithm to the dual of entropy regularized OT cuturi2013sinkhorn , researchers have conducted extensive studies in this area, including other types of regularization blondel2018smooth gasnikov2016efficient , acceleration guminov2021combination lin2022efficiency and numerical stability mandad2017variance . In xie2020fast , a proximal point algorithm (PPA) is considered to solve the discrete LP. The entropy inducing distance function is introduced to contruct the proximal term and each subproblem has the same formulation as the entropy regularized OT. This approach is found to stabilize the numerical computation while maintaining the efficiency of the Sinkhorn’s algorithm. In schmitzer2019stabilized , techniques such as log-domain stabilization, and epsilon scaling are discussed to further improve the performance of the Sinkhorn’s algorithm. The approach of iterative Bregman projection is discussed in benamou2015iterative . It is discovered that the Sinkhorn’s iteration is equivalent to a special case of this method. In pmlr-v84-genevay18a , a notion of Sinkhorn distance is proposed, which allows computable differentiation when serving as the loss function when training generative models. Sinkhorn’s algorithm can also be generalized to solve unbalanced OT in chizat2018scaling , and applied to solve OT over geometric domains in solomon2015convolutional . In huang2021riemannian , a Riemannian block coordinate descent method is applied to solve projection robust Wasserstein distance. The problem is to calculate Wasserstein distance after projecting the distributions onto lower-dimensional subspaces. The proposed approach employs entropy regularization, deterministic block coordinate descent, and techniques in Riemannian optimization. A domain decomposition method is considered in bonafini2021domain . Unlike our method, the decomposition is deterministic, focusing on one pair of distributions, and the authors try to tackle entropy-regularized OT.

Other works that significantly deviate from the entropy regularization framework include li2018computations , which computes the Schrödinger bridge problem (equivalent to OT with Fisher information regularization), and multiscale strategies such as gerber2017multiscale , liu2022multiscale and schmitzer2016sparse . In liu2022multiscale , the problem size is reduced using a multiscale strategy, and a semi-smooth Newton method is applied to solve smaller-scale subproblems. The RBCD method employed in this study is a regularization-free method. As a result, it avoids dealing with inaccurate solutions and numerical stability issues introduced by the regularization term. Furthermore, each subproblem in RBCD is a small-size LP, allowing for flexible resolution choices. Interior-point methods and simplex methods are also revisited and enhanced zanetti2023interior ; wijesinghe2023matrix ; natale2021computation ; gottschlich2014shortlist . In particular, the authors of zanetti2023interior propose an interior point inspired algorithm with reduced subproblems to address large-scale OT. The authors of facca2021fast consider an equivalent formulation of minimizing an energy functional to solve OT over graphs. Backward Euler (equivalent to the proximal point method) is used in the outer loop of the algorithm and Newton-Rapson is used in the inner loop. Other approaches include numerical solution of the Monge–Ampère equation benamou2014numerical ; benamou2016monotone , stochastic gradient descent to resolve OT in discrete, semi-discrete and continuous formulations genevay2016stochastic , and an alternating direction method of multipliers (ADMM) approach to solve the more general Wasserstein barycenter yang2021fast .

The rest of the paper is organized as follows. In Section 2, we review the basic idea of optimal transport and Wasserstein distance. In Section 3, we introduce the expected Gauss-Southwell-$q$ rule and a vanilla RBCD (RBCD₀) method for solving general LP problems. An inexact version of RBCD₀ is also discussed. In Section 4, we investigate the properties of the linear system in OT and propose several approaches to refine and accelerate the RBCD₀ method. In Section 5, preliminary numerical results are presented to demonstrate the performance of our proposed methods. Finally, concluding remarks are made in Section 6.

Motivated by the Gauss-Southwell-q rule introduced in tseng2009coordinate , we desire to select $\mathcal{I}_{k}$ such that

for some constant $v\in(0,1]$. Note that by (14), we have

where $x^{*}$ is an optimal solution of (10). Therefore, (12)-(17) imply that

(18) indicates that the gap of function value decays exponentially with rate $1-v$, as long as we choose $\mathcal{I}_{k}$ according to the Gauss-Southwell-q rule (16) at each iteration $k$. A trivial choice of $\mathcal{I}_{k}$ to satisfy (16) is $\mathcal{N}$ and $v=1$. However, this choice results in a potential large-scale subproblem in the BCD method (15), contradicting the purpose of using BCD. Instead, we should set an upper bound on the cardinality of $\mathcal{I}_{k}$, namely, a reasonable batch size to balance the computational effort in each iteration and convergence performance of BCD. Next, we discuss the existence of such an $\mathcal{I}_{k}$ given an upper bound $l$ on $|\mathcal{I}_{k}|$, which necessitates the following concept.

The following Theorem confirms the existence of such an $\mathcal{I}_{k}$ that satisfies (16), the proof of which follows closely to (tseng2009block, , Proposition 6.1).

However, it is not clear how to identify the set $\mathcal{I}$ described in Theorem 3.1 with little computational effort for a general $A$. Therefore, we introduced the following.

We introduce randomness in the selection of $\mathcal{I}_{k}$ to reduce the potential computation burden in identifying an $\mathcal{I}_{k}$ that satisfies (16). Consider an expected Gauss-Southwell-q rule:

where $v\in(0,1]$ is a constant, and $\mathcal{F}_{k}\triangleq\{x^{0},\ldots,x^{k}\}$ denotes the history of the algorithm. Therefore, using the notations of LP (10) and BCD method (15):

where $x^{*}$ is an optimal solution of (10). According to (polyak1987introduction, , Lemma 10, page 49), $c^{T}(x^{k}-x^{*})\to 0$ almost surely. Moreover, if we take expectations on both sides of (22),

i.e., the expectation of function value gap converges to $0$ exponentially with rate $1-v$.

Based on the expected Gauss-Southwell-$q$ rule, we formally propose a vanilla random block coordinate descent (RBCD₀) algorithm (Algorithm 1) to solve the LP (10). Specifically, we choose the working set $\mathcal{I}_{k}$ with full randomness, that is, randomly choose an index set of cardinality $l$ out of $\mathcal{N}$. Then with probability at least $\frac{1}{\binom{N}{l}}$, the index set will be the same as or cover the working set suggested by Theorem 3.1. As a result, (21) will be satisfied with $v\geq\frac{1}{\binom{N}{l}(N-l+1)}$.

Based on the previous discussions, Algorithm 1 generates a sequence $\{x^{k}\}$ such that the value of $c^{T}x^{k}$ converges to the optimal with probability 1. Moreover, the expectation of the optimality gap converges to $0$ exponentially. It is important to note that $\frac{1}{\binom{N}{l}(N-l+1)}$ is only a loose lower bound of $v$. This bound can become quite small when $N$ grows large due to the binomial coefficient $\binom{N}{l}$. However, in our numerical experiments (c.f. Sec. 5), this lower bound is rarely reached. In Section 4, we will discuss how to further improve this bound given the specific structure of the OT problem. Before that, we first investigate an inexact extension of RBCD₀.

For any $x$ and $\mathcal{I}\subseteq\mathcal{N}$, still denote $\mathcal{D}(x;\mathcal{I})$ and $q(x;\mathcal{I})$ as in (11) and (12). However, note that now $x$ does not have to be feasible. We consider the inexact version of Algorithm 1, where Step 2 in (1) is only approximately solved. We compute $d^{k}$ such that for any $d^{k}_{*}\in\mathcal{D}(x^{k};\mathcal{I}_{k})$,

where the inexactness sequences $\{\epsilon_{k}\}$ and $\{\hat{\epsilon}_{k}\}$ should be nonnegative. The inexact algorithm is described as follows.

Next, we analyze the sequence generated by Algorithm 1(a) and summarize the results in the following theorem.

A nonzero $d\in\mathbb{R}^{N}$ is an elementary vector of $\operatorname{null}(A)$ if $d\in\operatorname{null}(A)$ and there is no nonzero $d^{\prime}\in\operatorname{null}(A)$ that is conformal to $d$ and $\operatorname{supp}(d^{\prime})\neq\operatorname{supp}(d)$. According to the definition in (30), we say that a nonzero matrix $X$ is an elementary matrix of $\operatorname{null}(A)$ if $\operatorname{vec}(X)$ is an elementary vector of $\operatorname{null}(A)$. For simplicity, a matrix $M^{1}$ being conformal to $M^{2}$ means $\operatorname{vec}(M^{1})$ being conformal to $\operatorname{vec}(M^{2})$ for the rest of this paper.

Now we define a set $\mathcal{E}_{A}$:

First, we state a Lemma about $\mathcal{E}_{A}$, the proof of which is trivial and thus omitted.

Then we show every element in the null space of $A$ is related to a matrix in $\mathcal{E}_{A}$.