A fast and accurate splitting method for optimal transport: analysis and implementation
Authors: Vien V. Mai
Jacob Lindbäck
Mikael Johansson
Venue: ICLR 2022 Poster (2022)
Keywords: Optimal transport
Operator splitting
Douglas-Rachford
ADMM
GPUs

Abstract: We develop a fast and reliable method for solving large-scale optimal transport (OT) problems at an unprecedented combination of speed and accuracy. Built on the celebrated Douglas-Rachford splitting technique, our method tackles the original OT problem directly instead of solving an approximate regularized problem, as many state-of-the-art techniques do. This allows us to provide sparse transport plans and avoid numerical issues of methods that use entropic regularization. The algorithm has the same cost per iteration as the popular Sinkhorn method, and each iteration can be executed efficiently, in parallel. The proposed method enjoys an iteration complexity $O(1/\epsilon)$ compared to the best-known $O(1/\epsilon^2)$ of the Sinkhorn method. In addition, we establish a linear convergence rate for our formulation of the OT problem. We detail an efficient GPU implementation of the proposed method that maintains a primal-dual stopping criterion at no extra cost. Substantial experiments demonstrate the effectiveness of our method, both in terms of computation times and robustness.

Published as a conference paper at ICLR 2022
A FAST AND ACCURATE SPLITTING METHOD FOR OP- TIMAL TRANSPORT: ANALYSIS AND IMPLEMENTATION
Vien V. V.Mii
Jacob Lindbäck
Mikael Johansson
ABSTRACT
We develop a fast and reliable method for solving large-scale optimal transport (OT) problems at an unprecedented combination of speed and accuracy. Built on the celebrated Douglas-Rachford splitting technique, our method tackles the orig- inal OT problem directly instead of solving an approximate regularized problem. as many state-of-the of-the-art art techniques do. This allows us to provide sparse trans- port plans and avoid numerical issues of methods that use entropic regularization. The algorithm has the same cost per iteration as the popular Sinkhorn method, and each iteration can be executed efficiently, in parallel. The proposed method enjoys an iteration complexity O(1/e) compared to the best-known O(1/e²) of the Sinkhorn method. In addition, we establish a linear convergence rate for our formulation of the OT problem. We detail an efficient GPU implementation of the proposed method that maintains a primal-dual stopping criterion at no extra cost. Substantial experiments demonstrate the effectiveness of method, both in terms of computation times and robustness.
INTRODUCTION
We study the discrete optimal transport (OT) problem:
minimize
(C,
X20
subject to X1, = and X1 =9
where X ITT'" is the transport plan, C E 1R=*= is the cost matrix, and p E RT and E Ri are two discrete probability measures. OT has a rich history in mathematics and operations research dating back to at least the 18th century. By exploiting geometric properties of the underlying ground space, OT provides a powerful and flexible way to compare probability measures. It has quickly become central topic in machine learning and has found applications in countless learning tasks, including deep generative models (Arjovsky et al., 2017), domain adaptation (Courty et al., 2016), and inference of high-dimensional cell trajectories in genomics (Schiebinger al., 2019); refer to Peyré & Cuturi (2019) for a more comprehensive survey of OT theory and applications. However, the power of OT comes at the price of an enormous computational cost for determining the optimal transport plan. Standard methods for solving linear programs (LPs) suffer from a super-lineat time complexity in term of the problem size. Such methods are also very challenging parallelize on modern processing hardware. Therefore, there has been substantial research in developing new efficient methods for OT. This paper advances the state of the art this direction.
1.1 RELATED WORK
Below we review some of the topics most closely related to our work.
Sinkhorn method The Sinkhorn (SK) method (Cuturi, 2013) aims to solve an approximation of (1) in which the objective is replaced by regularized version of the form(C, (X). Here, (()) = - Ei Xıy log(X.,) is an entropy function and 17 > 0 a regularization parameter. The Sinkhorn method defines the quantity K = exp(-C) m) and repeats the following steps
U p/ (K(K- and Up g/(K" up).
"Ericsson AB, vien. ma www.ierrissoon.coo com KTH Royal Institute of Technology, (j nobac, mi kael jeekth.

---

Published as a conference paper at ICLR 2022
until lug O ((ex)-pl++]pe @ (K") ux) qll becomes small, then returns diag( K diag(es) The division (/) and multiplication (o) operators between two vectors are to be understood entry- wise. Each SK iteration is built from matrix-vectts multiplies and element-wise arithmetic opera- tions, and is hence readily parallelized on multi-core CPUs and GPUs. However, due to the entropic regularization, SK suffers from numerical issues and can struggle find even moderately accurate solutions. This problem is even more prevalent in GPU implementations as most modern GPUs are built and optimized for single-precision arithmetic (Kirk & Wen-Mei, 2016; Cheng et al., 2014). Substantial care is therefore needed to select an appropriate 7) that is small enough to provide a meaningful approximation, and large enough to avoid numerical issues. addition, the entropy term enforces dense solution, which can be undesirable when the optimal transportation plan itself of interest (Blondel et al., 2018). We mention that there has been substantial research in improving the performance of SK (Alaya et al., 2019; Altschuler et al., 2017: Blondel et al., 2018; 18; Lin et al., 2019). Most of these contributions improve certain aspects of SK: some result in more stable but much slower methods, while others allow to produce sparse solutions but at much higher cost per iteration. Some of these sophisticated changes make parallelization challenging due to the many branching conditions that they introduce.
Operator splitting solvers for general LP With relatively low per-iteration cost and the ability to exploit sparsity in the problem data, operator splitting methods such as Douglas- Rachford split- ting (Douglas & Rachford, 1956) and ADMM (Gabay & Mercier, 1976) have gained widespread popularity in large-scale optimization. Such algorithms can quickly produce solutions of moder- ate accuracy and are the engine of several successful first-order solvers (O' (O'doge et al., 2016; Stellato et al., 2020; Garstka et al., 2021). As OT can be cast as an LP, it can, in principle, be solved by these selittinn-basets solvers. However, there has not been much success reported in this context, probably due to the memory-bound nature of large-scale OTs. For OTs, the main bot- large
tleneck is not floating-point computations, but rather conecconu memory operations two-dimensional arrays. Even an innocent update like X + X C' is more expensive than the two matrix-vector multiplies in SK. To design high-performance splitting method, it is thus crucial to minimize the memory operations associated with large arrays. In addition, since most existing splitting solvers target general LPs, they often solve a linear system at each iteration, which is pro- hibitively expensive for many OT applications.
Convergence rates of DR for LP Many splitting methods, including Douglas- Rachford, are known to converge linearly under strong convexity (see c.g. Gisclsson & Boyd (2016)). Recently, it has been shown that algorithms based on DR/ADMM often enjoy similar convergence proper- ties also in the absence of strong convexity. For example, Liang al. (2017) derived local linear convergence guarantees under mild assumptions; Applegate et al. (2021) established global linear convergence for Primal-Dual methods for LPs using restart strategies; and Wang & Shroff (2017) es- tablished linear convergence for an ADMM- bbsed algorithm on general LPs. Yet, these frameworks quickly become intractable on large OT problems, due to the many memory operations required.
1.2 CONTRIBUTIONS
We demonstrate that DR splitting, when properly designed and implemented, can solve large-scale OT problems reliably and quickly, while retaining the excellent memory footprint and parallelization properties of the Sinkhorn method. Specifically, we make the following contributions:
We develop a DR splitting algorithm that solves the original OT directly, avoiding the numerical issues of SK and forgoing the need for solving linear systems of general solvers. We perform simplifications to eliminate variables so that the final algorithm can be executed with only one matrix variable, while maintaining the same degree of parallelization. Our method implicitly maintains primal-dual pair, which facilitates a simple evaluation of stopping criteria.
We derive an iteration complexity O(1/e) for our method. This is a significant improvement over the best-known estimate O(1/e²) for the Sinkhorn method (cf. Lin et al. (2019). We also provide a global linear convergence rate that holds independently of the initialization, despite the absence of strong convexity in the OT problem.
We detail efficient GPU implementation that fuses many immediate steps into and per- forms several on mn-the-fll fly reductions between a read and write of the matrix variable. We also show how pril stopping criterion can be evaluated extra cost. The implementation is
2

---

Published as a conference paper at ICLR 2022
available open source and gives practitioners fast and robust OT solver also for applications where regularized OT not suitable.
As by-product of solving the original OT problem, our approximate solution is guaranteed to be sparse. Indeed, it is known that DR can identify an optimal support in a finite time, even before reaching solution lutzeler & Malick (2020). To avoid cluttering the presentation, we focus on the primal OT, but note that the approach applies equally well to the dual form. Moreover, our implementation readily extended to other splitting schemes (e.g. Chambolle & Pock (2011)).
2 BACKGROUND
Notation For any 2, E R", (x, is the Euclidean inner product of x and y, and 1-ll2 denotes the f2-norm. For matrices X,Y E Rmxx, (X, Y) = tr(X' y) denotes their inner product and ||-Il- (..-) is the the induced Frobenius norm. We use ||-Il to indicate either ||-1|g or ||-llp- For a closed and convex set X, the distance and the projection map are given by dist(r, X) millze ||= xi| and Px(x)) argmingex || =||, respectively. Further, we denote the proximal operator of a closed convex function f by prox, of (x) = argmin, f(z) + 2p
1 I| ==1===1>1000 > 0. The
Euclidean projection of R onto the nonnegative orthant denoted by [xr] max( 0), and
An {i E RT > Ij = 1} is the (n - 1) dimensional probability simplex.
OT and optimality conditions Lete = 1n and f = 1m. and consider the linear mapping
A : mx + mm +n X + (Xe, XT)).
and its adjoint A Rm+ mmxn xv. ye? + fxJ The projection the range of A a -)), where a = ( y - -eTx) (xT((mmm + n) (Bauschke et al., 2021).
is Prana((V) x)) = (y. x)
With (p.q) m Problem (1) can be written as linear program on the form
minimize XERmXn
X)
subject to A(X) b X20 0.
Let (1,)) E mm++ be the dual variable associated the affine constraint in (2). Then, using the definition of A", the dual problem of (2), or equivalently of (1), reads
maximize p pu
subject to pe + ffvTSC.
(3)
OT is bounded program and always admits feasible solution. It is thus guaranteed to have an optimal solution, and the optimality conditions can be summarized as follows.
Proposition! 1. A pair X and (ps, v) are primal and dual optimal if and only if: (i) XT5=
q, X 2 (ii) ue + fv C O, (iii) (X,C pe - ff v)
These conditions mean that: (1) X is primal feasible; (ii) μ,, dual feasible, and (iii) the duality gap is zero (C, = p_ μ + q v, or equivalently, complementary slackness holds. Thus, solving the OT problem amounts (approximately) finding such primal-dual pair. To this end, we rely on the celebrated Bouglas-Rachfon splitting method (Lions & Mercier, 1979; Eckstein & Bertsekas, 1992; Fukushima, 1996; Bauschke & Combettes, 2017).
Nouglas-Rachford splitting Consider composite convex optimization problems on the form
minimize f(x) + g(x), =ER"
44)
where and g proper closed and convex functions. To solve Problem (4), the DR splitting algorithm starts from yo E R" and repeats the following steps:
Dh+1 prox,f pf (ya), Zx+1 = prox (2xx+1 yx), Mk++ = yk + 2k+1 2k+1
(5)
where p > 0 is a penalty parameter; prox,5 (-) and proxpy (-) are the proximal operator of f and g. respectively, Procedure (5) can be viewed as a fixed-point iteration 1 = Tyy) for the mapping
T(y) = y + prox pg (2prox, os (v) v) - prox,s (v).
(6)
3

---

Published as a conference paper at ICLR 2022
The DR iterations (5) can also be derived from the ADMM method applied to equivalent problem to (4) (see, Appendix C). Indeed, they are both special instances of the classical proximal point method in Rockafellar (1976). As for convergence, Lions & Mercier (1979) showed that T is firmly-nonexpansivee mapping, from which they obtained convergence of yk- Moreover, the sequence Jy is guaranteed to converge to minimizer of f (assuming a minimizer exists) for any choice of > O. In particular, we have the following general convergence result, whose proof can be found in Bauschke & Combettes (2017, Corollary 28.3).
Lemma 2.1. Consider the composite problem (4) and its Fenchel-Rockafelee dual defined as maximize uER"
(f" - (u).
Let P and D be the sets of solutions to the primal (4) and dual (7) problems, respectively, Let Tk. Mke and zy be generated by procedure (5) and let ug == (2k-1 Tx)/p. Then, there exists y" R= such that yk + Setting x² proxpf (y") and u² (y* z²)/p. then it holds that (i)
x E P* and u² E D (ii) 22 2 + O, I + x= and 2 + x;; (iii) up + u²
3 DOUGLAS-RACHIEE
SPLITTING FOR OPTIMAL TRANSPORT
To efficiently apply DR to OT, we need to specify the functions f and g as well as how evaluate their proximal operations. We begin by introducing result for computing the projection onto the set of real-valued matrices with prescribed row and column sums (Romero, 1990; Bauschke al., 2021).
Lemma 3.1. Let e = 1, and f = 1m Letp E Am and 9 E An Then, the set X defined by:
X : (X m= Xe = p and = = =q}
is non-emply, and for every given X E BRmxn, we have
Px(X)= x X ((Xe p) -7se 1 (S(XTS -(fe))
where 7 = 5T (Xe p) (m n) = e (xT 4) (m + n).
The lemma follows immediately from Bauschke et al. (2021, Corollary 3.4) and the fact that (p, = PranA((p, q)). implies that (X can be carried out by basic linear algebra operations, such as matrix-vector multiplies and rank-one updates, that can be effectively parallelized.
3.1 ALGORITHM DERIVATION
Our algorithm is based on re-writing (2) on the standard form for DR-splitting (4) using carefully selected f and g that ensures rapid convergence of the iterates and an efficient execution of the iterations. particular, we propose to select and as follows:
f(X) = XC, + 1g (X) and g(X) = [o):A
(8)
By Lemma 3.1, readily have the explicit formula for the proximal operator of g. namely, prox pg (-) = P-))). The proximal operator f can also be evaluated explicitly as:
prox, (X) P mx (X pC) [X pCpC+
The Douglas-Rachford splitting applied to this formulation of the problem then reads:
X [Yr - pCC+. Z = (2X - Yk), YA+1 = Y2 + Z X (9)
Despite their apparent simplicity, the updates in (9) are inefficient to execute in practice due to the many memory operations needed to operate on the large arrays Xx. Yg, Zy and C. To reduce the memory access, will now perform several simplifications to eliminate variables from (9). The resulting algorithm can be executed with only one matrix variable while maintaining the same degree of parallelization.
We first note that the linearity of (x) allows us to eliminate Z, yielding the Yuupdate
Ye+1 = Xx+1 n (2X++1e YAe Tmf)eT mn 1(22 Yg T f mm))
4

---

Published as a conference paper at ICLR 2022
Algorithm Soaugla-RRacffort Splitting for Optimal Transport (DROT)
Input: OT(C, 9). initial point Xo, penalty parameter p
1: Φo - 0, :0
2: an Xoe p. bo
X S 4, 00 = /aao(mm + n)
3: for = 2.1.2.... do
4: X [x ØNE' + fPp -pC]+
5: Tk+1 Xx+1E p. 8k41 = Xxt1 - 9, Bx+1 = S TT+xiimmm+
6: Φk+1 = (ay: 2r&+1 + (23x+1 - f) /n
7: (k+11 (bx- 28k+1 (23 as)e) /m
8: 0k+1 = aę rk+1, bx+1 = bk- 8kt1, ax+1 = 0k - 3k+1
9: end for
Output: XK
f where "Yk = f (2Xx+1e Yye - (pe--))/oe /m +n) = E (2x fYY 9) (m + n). We also
define the following quantities that capture how Yú: affects the update of Y k+1
ak = Yye= p. by Yx
as TTat/a//n am/(mm + n) = e b + n))
Similarly, for Xu. we let:
Tk = xxe- P.
8 = x 8g f + = T sx/(m n).
Recall that the pair (Fks 9k) represents the primal residual at Xg. Now, the preceding update can be written compactly as
Yk+1 = Xk+1 + oxt1e + fuk+11
(10)
where
- 6x+1 = =(a; - -2rx+1 + (22)+110x)fff1 /n
pk++1 (bx- 2sk+1 (28k+1 - ax)e) /m.
We can see that the Yuupdate can be implemented using 4 matrix-vector multiples (for computing ag, ba, 1,+1, 8k+1), 1), followed by 2 rank-one updates. As a rank-one update requires read from an input matrix and write to an output matrix, typically twice as costly as a matrix-vector multiply (which only writes the output to vector). Thus, it would involve 8 memory operations of large arrays, which is still significant.
Next, show that the Yupddate can be removed too. Noticing that updating Yx+1 from Yx and XA-1 does not need the full matrix Yg, but only the ability to compute ag and by- This allows us to use single physical memory array represent both the sequences X and Yg. Suppose that overwrite the matrix XA-1 as:
X X QRtieT f
then after the two rank-one updates, the X-array holds the value of Ys+1- We can access the actual X-value again in the next update, which now reads: X8+2 + [X++1 - pC) It thus remains to show that ag and by can be computed efficiently. By multiplying both sides of (10) by and subtracting the result from p. we obtain
1 ---xxtte--- de-1e Te e (xT+1e)).
Since e = n and (hx 28x+1) = (m + n) 28++1), holds that (pI+1)) = (ox 2f. We also have le e ag 2 + (23 - ox)f. Therefore, we end up with an extremely simple recursive form for updating 04:
ax+1 = 0k Tk+1
Similarly, we have b = by 9k+1 and ax+1 = ag SktT . In summary. the DROT method can be implemented with single matrix variable summarized in Algorithm 1.
5

---

Published as a conference paper at ICLR 2022
Stopping criteria It interesting to note that while DROT directly solves to the primal prob- lem (1), it maintains a pair vectors that implicitly plays the role of the dual variables μ and v in the dual problem (3). To get a feel for this, note that the optimality conditions in Proposition are equivalent to the existence of a pair X and (n² v²) such that
ex f) (p.g) and X = |x* + p= + fv
Here, the later condition encodes the nonnegativity constraint, the dual feasibility, and the zero duality gap. The result follows by invoking Deutsch (2001, Theorem 5.6(i)) with the convex cone X* E C and 2 , + + T - C' C Now, compared to Step 4 in DROT, the second condition above suggests that Φ p and px/p are such implicit variables. To see why this indeed the case, let U = (Ya-1 xx) Then it is easy to verify that
C == R4
(Z₁+ Xx) (Xx - Yk-1 + oxe + ff/))/p = -Ux + (ox/p)e² + f(px//)).
By Lemma 2.1, we have Zk Xx and UK - U E mmxx by which it follows that
(dx/p)e + [((xx/)) -
Finally, by evaluating the conjugate functions f and g* in Lemma 2.1, it can be shown that U must have the form u + fu² where μ E R and V^ E IR satisfying "²e" + fv < C; see Appendix D for details. With such primal- pairs at our disposal, we can now explicitly evaluate their distance to the set of solutions laid out in Proposition I by considering: Tprimal (|]r=|1 || 113/ rdual = Il[(@x/p)e" + f(pa/p) C]+ and gap = (C,Xa) (p ox + skk 4)/pl- As problem (1) is feasible and bounded, Lemma 2. and strong duality guarantees that all the three terms will converge zero. Thus, we terminate DROT when these become smaller than some user-specified tolerances.
3.2 CONVERGENCE RATES
In this section, state the main convergence results of the paper, namely sublinear and linear rate of convergence of the given splitting algorithm. In order to establish the sublinear rate, we need the following function:
V(X,
= f(x) + (U,Z
which is defined for X R Z E X and U Rmxx, We can now state the first result.
Theorem 1. Let Xx, Yk, Zg be the generated by procedure (9). Then, for any X Rmxn Z E X,
and Y Ruxn we have
V(X++1, Ze42, (Y -)/p) V(,, XK+1)/P) S (11Y- 1| r11)) ((2p).
Furthermore, let Y be a fixed-poin of T in (6) and let X be a solution of (1) defined from Y in the manner of Lemma 1. Then, it holds that
(c,Xx) (C, - (I261 /(2p) + + 1Yo - 1x-1/p)) /k
12 - Xxll = \ / - Y*1l//,
where X, Et X,/k and Z Σx1 /i/k.
The theorem implies that one can compute a solution satisfying (C, xx (C, -C, S c O(1/)) iterations. This is significant improvement over the best-known iteration complexity 0(1/e²) of the Sinkhorn method (cf. Lin et al. (2019)). Note that the linearity of (C.-) allows to update the scalar value (C, X²) recursively without ever needing to form the ergodic sequence Xg. Yet. in terms of rate, this result is still conservative, as the next theorem shows.
Theorem 2. Let X and Yn be generated by (9). Let G* be the set of fixed points of T in (6) and let X* be the set of primal solutions to (1). Then, {Yr} is bounded, |Y'll M for all k, and
dist (Ys, g*) ) S (Yst(YY. (Y..g))))) x b and dist xx,,2*))
X²) & dist 9*) x k-1
where r = v/ +I < c= y(1 + p(llell + IIfID) 2 1, and y 05- (1 + p (M + 1)). Here, 0s is prohlem-dependent constant characterized by the primal- dual solution sets only.

---

Published as a conference paper at ICLR 2022
WW13-16. sdueres
15.00
Gubilreren P-ll
Thrsud regiin be
Figure 1: Left: Logical view of the main kernel. To expose sufficient parallelism GPU, we organize threads into 2D grid of blocks (3 x 2 in the figure), which allows several threads per row. The threads are then grouped in ID blocks (shown in red) along the columns of X. This ensures that global memory access is aligned and coalesced to maximize bandwidth utilization. We use the parameter work-size ws to indicate how many elements of row each thread should handle. For simplicity, this parameter represents multiples of the block size bs. Each arrow denotes the activity a single thread in thread block. Memory storage assumed to be column column-major Right: Activity of a normal working block, which handles a submatrix of size bs x (ws-bo) bs).
This means that an e-optimal solution can be computed in (log 1/e) iterations. However, it is, in general, difficult to estimate the convergence factor r, and it may in the worst case be close to one. In such settings, the sublinear rate will typically dominate for the first iterations. either case, DROT always satisfies the better of the two bounds at each k.
3.3 IMPLEMENTATION
In this section, we detail our implementation of DROT to exploit parallel processing on GPUs. We review only the most basic concepts of GPU programming necessary to describe our kernel and refer to Kirk & Wen-Mei (2016); Cheng et al. (2014) for comprehensive treatments.
Thread hierarchy When a kernel function is launched, large number of threads are generated to execute its statements. These threads are organized into two-level hierarchy. A grid contains mul- tiple blocks and a block contains multiple threads. Each block is scheduled to one of the streaming multiprocessors (SMs) the GPU concurrently While all threads in a thread block run logically in parallel, not all of them can run physically at the same time. As result, different threads in thread block may progress a different pace. Once thread block is scheduled to an SM, its threads are further partitioned into warps. A warp consists of 32 consecutive threads that execute the same instruction the same time. Each thread has its own instruction address counter and register state, and carries out the current instruction on its own data.
sequentially, depending
available hardware.
Memory hierarchy Registers and shared memory ("on-chip") are the fastest memory spaces on a GPU. Registers are private to each thread, while shared memory is visible to all threads in the same thread block. An automatic variable declared memory is programmable, and users have full control over when data is moved into or evicted from the shared memory. It enables block-level communication, facilitates reuse of on-chip data, and can greatly reduce the global memory access of kernel. However, there are typically only a couple dozen registers per thread and a couple of kBs shared memory per thread block. The largest memory
kernel generally stored in a register. Shared
GPU card is global memory, physically separated from the compute chip ("off- chip"). All threads can access global memory, but its latency is much higher, typically hundreds of clock cycles.! ! Therefore, minimizing global memory transactions is vital to performrree kernel. When all threads of a warp execute load (store) instruction that access consecutive memory locations, these will be coalesced into as few transactions as possible. For example, if they access consecutive 4-byte words such as float32 values, four coalesced 32-byte transactions will service that memory access.
Before proceeding further, we state the main result of the section.
Claim 3.1. On average, an iteration of DROT, including all the stopping criteria, can be done using 2.5 memory operations of m x n arrays. In particular, this includes one read from and write to X, and one read from C in every other iteration.
'https://dccs nwi didiia cctm://uuaa/cuaaa--orrrrraaaaaaa wwwpraamminnggiiiin html
7

---

Published as a conference paper at ICLR 2022
Main kernel Steps 4-5 and the stopping criteria are the main computational components of DROT, since they involve matrix operations. We will design an efficient kernel that: (i) updates X to Xx-1, (ii) computes uk+1 == X e and v == X (iii) evaluates (++1) in the duality gap expression. The kernel fuses many immediate steps into one and performs several on-the-fly reductions while updaring Xs. thereby minimizing global memory access.
Our kernel is designed so that each thread block handles a submatrix x of size bs x (ws)bs), except for the corner blocks which will have fewer rows and/or columns. Since all threads in block need the same values of yp. it is best to read these into shared memory once per block and then let threads access them from there. We, therefore, divide yp into chunks of the block size and set up loop to let the threads collaborate in reading chunks in a coalesced fashion into shared memory. Since each thread works on a single row, it accesses the same element of Φ throughout, and we can thus load and store that value directly in a register. These allow maximum reuse of yp and Φ
In the j-th step, the working block loads column j of x to the chip a coalesced fashion. Each thread i E {0, 1,..., bs - 1} uses the loaded zij to compute and store Zij in register:
z.j = --maax (Tag + Φ; + [P] - pCkj: 0),
where c is the corresponding submatrix of C. As sum reduction is order- inder-independent, it can be done locally at various levels, and local results can then be combined to produce the final value. We, therefore, reuse 2rj and perform several such partial reductions. First, to compute the local value for 12 at column j. thread simply adds #ty to a running sum kept in a register. The reduction leading to (C, K++1) can be done in the same way. The vertical reduction to compute 0k+1 is more challenging as coordination between threads is needed. We rely on warp-level reduction in which the data exchange is performed between registers. This way, we can also leverage efficient CUDA's built-in functions for collective communication at warp-level. When done, the first thread in warp holds the total reduced value and simply adds it atomically to proper coordinate of Uş-1 Finally, all threads write o to the global memory. The process repeated by moving to the next column.
When the block finishes processing the last column of the submatrix æ, each thread adds its running sum atomically to a proper coordinate of ug+1- They then perform one vertical (warp-level) reduc- tion to collect their private sums to obtain the partial cost value (c, ær). When all the submatrices have been processed, we are granted all the quantities described in (iii, as desired. It is essen- tial to notice that each entry of æ is only read and written once during this process. To conclude Claim 3.1, we note that one can skip the load of C' in every other iteration. Indeed, if in iteration k. instead of writing Xk+1, one writes the value of X - PC, then the next update is simply X [X*+1 + ØR+iE + f Finally, all the remaining updates in DROT only involve vector and scalar operations, and can thus be finished off with simple auxiliary kernel.
4 EXPERIMENTAL RESULTS
In this section, we perform experiments validate our method and demonstrate its efficiency both in terms of accuracy and speed. We focus on comparisons with the Sinkhorn method, as implemented in the POT toolbox², due to its minimal per-iteration cost and its publicly available GPU implementation. All runtime tests were carried out on an NVIDIA Tesla T4 GPU with 16GB of global memory. The CUDA C++ implementation of DROT is open source and available at https: //github. com/ vit.maiiiiaiiimmmmmmmmmmmmmm drot.
We consider six instances of SK and its log-domain variant, called SKI-SK6 and SK-Log-I-SK-
Log-6, corresponding to 77 10 1001 3,- 10 3,000000 3,0 x 10 10-2,10-7 2,101,, in that order. Given m and n, we generate source and target samples as 2g and 2t, whose entrics are drawn from a 2D Gaussian distribution with randomized means and covariances (Mg, Σ.)) and (14, E.). Here, /4/ R. has normal distributed entries, and Σ. AAAA. where A, E 2x2 is a matrix with random entries in [0, 1]; Me E 2 has entries formed from (5 o,) for o, > O, and Σ, E R2x2; is generated similarly to Eg. Given xa and Tr, the cost matrix C' represents pair-wise squared distance between samples: C1y=1|=a - 2, [ for i E {1...., m and j {l.....n? n and is normalized to have ||C'| = 1.
The marginals and are set to p- 1m/ m and q = 1,/n. DROT always starts at X, = pq
"nttps://devel nvi www.ia..cioooiiiii....o wtt.d/a....o.ccaa//eooooooooooco ing cada--cda-aaa level devel-primi ves/
8

---

Published as a conference paper at ICLR 2022
EROOT
ERST
33
(a) Float 32, Ot
(b) Float 64, O
(c) Float 64, o5 10
Figure 2: The percentage of problems solved up to various accuracy levels for o = 5, 10.
Performance profile For each method, we evaluate for each e > 0 the percentage of experiments for which f(x)) fx I e, where f(X*) is the optimal value and f(Xx) is the objective value at termination. An algorithm is terminated as soon as the constraint violation goes below 10-4 or 1000 iterations have been executed. Figure 2 depicts the fraction of problems that are successfully solved up to an accuracy level e given on the z-axis. Each subplot shows the result on 100 random problems with m n = 512 and p = 2/(m + n). We can see that DROT is consistently more accurate and robust. It reinforces that substantial care is needed select the right 7 for SK. The method is extremely vulnerable SK instance that seems to work in one setting can into numerical issues in another. For example, by slightly changing a statistical properties of the underlying data, nearly 40% of the problems in Fig. 2(c) cannot be solved by SK2 due to numerical errors, even though it works well in Fig. 2(b).
and even in double-precision, an
Runtime As we strive for the excellent per-iteration cost of SK, this work would be incomplete if we do not compare these. Figure 3(a) shows the median of the per-iteration runtime and the 95% confidence interval, as a function of the dimensions. Here, m and n range from 100 to 20000 and each plot is obtained by performing 10 runs; in each run, the per-iteration runtime is averaged over 100 iterations. Since evaluating the termination criterion in SK is very expensive, we follow the POT default and only do once every 10 iterations. The result confirms the efficiency of our kernel, and for very large problems the per-iteration runtimes of the two methods are almost identical. is also clear that DROT and SK can execute iteration much faster than the log-domain SK implementation. Finally, Figs 3(b)-cc) show the median times required for the total error (the sum of function gap and constraint violation) to reach different e values. Since SK struggles to find even moderately accurate solutions in single precision, we drop these from the remaining runtime comparisons. Note also that by design, all methods start from different initial points.* We can see that DROT can quickly and consistently find good approximates, while SK-Log is significantly slower for small y and may not attain the desired accuracy with higher 17.
15-4
19
(a) Runtime per iteration
(b) Time to E = 0.01
(c) Time to E 0.005
Figure Wall-clock runtime performance versus the dimensions m = n for a;
The experiments confirm the strength of DROT both in speed and robustness: (i) each DROT itera- tion can be executed efficiently in SK, and (ii) it can find high accuracy solutions as a well-tuned log-domain SK does, but at a much faster speed.
*The SK methods have their matrices K = exp(-C/n) and Xo diag (No) K diag(vv). For DROT, since X = [X + Øue f& pC]+, where Xo = pq 1 1/(mn). ||C 20 = 1. = Po/ Po/(m + n). do : = 0. the term inside [+]+ is of the order O(1((mn)) pm/(m + +)) & 0. This makes the first few iterations of DROT identically zeros. To keep the canonical Xo, we just simply set po to 0.5/ log(m) to reduce the warm-up period as increases, which is certainly suboptimal for DROT performance.

---

Published as a conference paper at ICLR 2022
ACKNOWLEDGEMENT
This work was supported in part by the Knut and Alice Wallenberg Foundation, the Swedish Re- search Council and the Swedish Foundation for Strategic Research, and the Wallenberg AI, Au- tonomous Systems and Software Program (WASP). The computations were enabled by resources provided by the Swedish National Infrastructure for Computing (SNIC), partially funded by Swedish Research Council through grant agreement no. 2018-05973.
REFERENCES
Mokhtar Alaya, Maxime Berar, Gilles Gasso, and Alain Rakotomamoniy. Screening Sinkhorn al- gorithm for regularized optimal transport. In Advances in Neural Information Processing Systems, volume 32, 2019.
Jason Altschuler, Jonathan Weed, and Philippe Rigollet. Near-linear time approximation algorithms for optimal transport via Sinkhorn iteration. In Advances in Neural Information Processing Sys- tems, pp. 1964-199 1974, 2017.
David Applegate, Oliver Hinder, Haihao Lu, and Miles Lubin. Faster first-order primal-dual methods for linear programming using restarts and sharpness. arXiv preprint arXiv:2105.122/5, 2021.
Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein generative adversarial networks. In International conference on machine learning, pp. 214-223. PMLR, 2017.
Heinz H Bauschke and Patrick L Combettes. Convex analysis and monotone operator theory in Hilbert spaces. Springer, 2nd edition, 2017.
Heinz H Bauschke, Shambhavi Singh, and Xianfu Wang. Projecting onto rectangular matrices with prescribed row and column sums. arXiv preprint arXiv:2105. 12222, 2021.
Amir Beck. First-order methods in optimization, volume 25. SIAM, 2017.
Mathieu Blondel, Vivien Seguy, and Antoine Rolet. Smooth and sparse optimal transport. In Inter- national Conference Artificial Intelligence and Statistics, pp. 880-889, PMLR, 2018.
Antonin Chambolle and Thomas Pock. A first-order primal-dual algorithm for convex problems with applications to imaging. Journal of mathematical imaging and vision, 40( 40(1):120-145, 2011.
John Cheng, Max Grossman, and Ty McKercher. Professional CUDA C programming. John Wiley & Sons, 2014.
Nicolas Courty, Rémi Flamary, Devis Tuia, and Alain Rakotomamoniy. Optimal transport for do- main adaptation. IEEE transactions on pattern analysis and machine intelligence, 39(9):1853- 1865, 2016.
Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimml transport. Advances in neural information processing systems, 26:2292-2308 2013.
Frank Deutsch. Best approximation inner product spaces, volume Springer, 2001.
Jim Douglas and Henry H Rachford. On the numerical solution of heat conduction problems in two and three space variables. Transactions of the American mathematical Society, 82(2):421-45 1956.
Jonathan Eckstein and Dimitri P Bertsekas. On Douhlas-Rachford splitting method and the proximal point algorithm for maximal monotone operators. Mathematical Programming, 55(1): 293-318, 1992.
Masao Fukushima. The primal Douglas-Rachfon splitting algorithm for class of monotone map- pings with application to the traffic equilibrium problem. Mathematical Programming, 72(1): 1-15, 1996.
10

---

Published as a conference paper at ICLR 2022
Daniel Gabay and Bertrand Mercier. A dual algorithm for the solution of nonlinear variational problems via finite element approximation. Computers & mathematics with applications, 2(1): 17-40, 1976.
Michael Garstka, Mark Cannon, and Paul Goulart. COSMO: convex conic problems. Journal of Optimization Theory and Applications, pp- 1-32, 2021.
conic operator splitting method for
Pontus Giselsson and Stephen Boyd. Linear convergence and metric selection for Douglas Rachford splitting and admm. IEEE Transactions on Automatic Control, 62(2):532-544, 2016.
Bingsheng He and Xiaoming Yuan. On the O(1/n) convergence rate of the Douglas-Rachin Rachford alternating direction method. SIAM Journal on Numerical Analysis, 50(2):700-709, 2012
Franck lutzeler and Jérôme Malick. Nonsmoothness machine learning: specific structure, proxi- mal identification, and applications. Sei-Valued and Variational Analysis, 28(4):661-678, 2020.
David Kirk and W Hwu Wen-Mei. Programming massively parallel processors: hands-on approach. Morgan kaufmann, 2016.
Wu Li. Sharp Lipschitz constants for basic optimal solutions and basic feasible solutions of linear programs. SIAM journal on control and optimization, 32(1):140-153 1994.
Jingwei Liang, Jalal Fadili, and Gabriel Peyré. Local convergence properties of Douglas-Rachford and alternating direction method of multipliers. Journal of Optimization Theory and Applications, 172(3):874-91 913, 2017.
Tianyi Lin, Nhat Ho, and Michael Jordan. On efficient optimal transport: An analysis of greedy and accelerated mirror descent algorithms. In International Conference on Machine Learning, pp. 3982-3991, PMLR, 2019.
Pierre-Louss Lions and Bertrand Mercier. Splitting algorithms for the sum of two nonlinear opera- tors. SIAM Journal on Numerical Analysis, 16(6):964-979, 1979.
Brendan O'donoghue, Eric Chu, Neal Parikh, and Stephen Boyd. Conic optimization via operator splitting and homogeneous self-dual embedding. Journal of Optimization Theory and Applica- tions, 169(3):1042- 1068, 2016.
Gabriel Peyré and Marco Cuturi. Computational optimal transport. Foundations and Trends@ in Machine Learning, 1 (56-6))355-6677 2019.
Tyrrell Rockafellar. Monotone operators and the proximal point algorithm. SIAM journal on control and optimization, 14(5):877-898, 1976.
David Romero. Easy transpoortation-like problems onk-dimensional arrays. Journal of optimization theory and applicaaioon,, 66(1): 66(1):137-147, 1990.
Geoffrey Schiebinger, Jian Shu, Marcin Tabaka, Brian Cleary, Vidya Subramanian, Aryeh Solomon, Joshua Gould, Siyan Liu, Stacie Lin, Peter Berube, et al. Optimal-transport analysis of single-cell gene expression identifies developmental trajectories in reprogramming, Cell, 176(4):928-443, 2019.
Bartolomeo Stellato, Goran Banjac, Paul Goulart, Alberto Bemporad, and Stephen Boyd. OSQP: An operator splitting solver for quadratic programs. Mathematical Programming Computation, pp. 1-36, 2020.
Sinong Wang and Ness Shroff. A new alternating direction method for linear programming. In Proceedings of the 31st International Conference on Neural Information Processing Systems, pp- 1479-1487, 2017.
11

---

Published as a conference paper at ICLR 2022
A EXTRA EXPERIMENTS
This section complements the experimental results in the main paper with further experiments for DROT, SK and SK-Log in various settings.
Performance profile for MNIST We follow the construction in Cuturi (2013) and generate the cost matrix C' E IR 784 the normalized matrix of squared Euclidean distances between 28 28 bins in the grid. For each pair of source and target digits, we convert an image into vector of intensity of 28 28 and normalize it to form the probability vectors p and q. respectively. Similarly to Figure 2, Figure shows the fraction of problems that are successfully solved up to an accuracy level e given on the x-axis. Here, each subplot is obtained by comparing 300 random pairs of images.
DROT
BROT
0.5
0.3
Perfrmatee oke
0.6
0.6
0.2
0.2
10-1
10a
10th
(a) Float 32
(b) Float 64
Figure 4: The percentage of problems solved up to various accuracy levels for MNIST.
Performance profile for SK-Log Figure 5 shows the fraction of problems that successfully solved by DROT and the log-domain implementation of SK. We can see that even with a simple stepsize rule, DROT achieves similar performance of well-tuned log-domain SK method.
1.0
0.8
Performance ratlo
DROT
0.4
SK-Log-1
SK-Log-2
SK-Log-3
0.2
5K-Log-4
5K-Log-5
0.0
SK-Log-6
10
10
10
10-1
Accuracy
Figure 5: The percentage of problems solved up to various accuracy for a, = 5 and p= 2/22mm +n).
Single runs performance Figure 6 depicts the actual performance of different algorithms for ran-
dom OT instances used to produce the performance profile in Figure In addition, Figure 7 high- lights the linear convergence phenomenon predicted by our theory.
Sparsity of transportation By design, DROT efficiently finds sparse transport plans. illustrate
this, we apply DROT to color transfer problem between two images (see Blondel et al. (2018)). By doing so, obtain highly sparse plan (approximately 99% zeros), and in turn, high-quality
12

---

Published as a conference paper at ICLR 2022
100
109
10"4
10.2
Nuncston Gae
10'7
Punetton Cam
10-2
10-1
10.1
DROT
10-4
10
5K66
5KG
200
1900
Iteration
(a) Float 32, 0;
Float 64, oe
Figure 6: Objective suboptimality versus iteration for random OT instances with = 12 512.
10a
DAXT
10
DROT
10
10-
10'1-
Funcion Gap
10
Rundier 63p
10"7
10-
10:
10-13
1000
2000
3000
1000
5000
2000
4000
6300 Iterattoon
8000
10ice
12000
Reration
(a) 11 32,
= 1.25
(b) 64, 1.33
Figure 7: Objective suboptimality versus iteration for random OT instances with d = 1.
artificial image, visualized in Figure 8. In the experiment, we quantize each image with KMeans to reduce the number distinct colors to 750. We subsequently use DROT to estimate an optimal color transfer between color distributions of the two images.
Source
1000 DROT Iterations
Target
Figure 8: Color transfer via DROT: The left-most image a KMeans compressed source image (750 centroids), the right-most is compressed target image (also obtained via 750 KMeans centroids). The middle panel displays an artificial image generated by mapping the pixel values of each centroid in the source weighted mean of the target centroid. The weights are determined by the sparse transportation plan computed via DROT.
13

---

Published as a conference paper at ICLR 2022
B PROOFS OF CONVERGENCE RATES
In order to establish the convergence rates for DROT, we first collect some useful results associated with the mapping that underpins the DR splitting:
T(y) = y + prox pg (2prox,, of (v) - y) - prox pf (v)
DR corresponds to finding a fixed point to T. i.c. a point y" y² = Ty"). or equivalently finding a point in the nullspace of G, where G(yy) y- Ty). The operator T,, and thereby also G, are are firmly non-expansive (see e.g. Lions & Mercier (1979)). That is, for all y, E dom(T):
(T(y) - T(1),, y 3/) 2 (T(y) T(1)
(G(y) - G(y)), y - -yy 2 [G( G()))ll
(11)
Let G* be the set of all fixed points to T, ie. G {y E dom(T)) G(y) 0}. Then (11) implies the following bound:
Lemma B.1. Let y E e and yk+1 (yk)), it holds that
lluk41 112 1lve v²12 11 yn
Proof. We have
11 - 12
11 yk yk y/"
1 y 21 + 2 (yx+1 - yk 3/ - 4²) + 11 - y/"
11 3/ 112 2 (G((x) - -Gy(()) yk y²) + 11 y 11?
< 11 y/ 112 21IG((m) - -II + Il 1²11²
1l 112 11 3all22.
Corollary B.1. Lemma B.I implies that
|| y/s $ 11ye- y'll
(12a)
(Jist())=1, & (dist g*))² -11 1xxl1
(12b)
Proof. The first inequality follows directly from the non-negativity of the left-hand side of Lemma B.I. The latter inequality follows from:
(dist(y--1.9") ||y++1 Ps- (y+++))l S 11 Ps- (1x)
Applying Lemma B.1 to the right-hand side with y" Pg- (yk) yields the desired result.
B.1 PROOF OF THEOREM I
Since DR is equivalent to ADMM up to changes of variables and swapping of the iteration order, expected that DR can attain the ergodic rate (11kk) of ADMM, both for the objective gap and the constraint violation (He & Yuan, 2012; Beck, 2017). However, mapping the convergence proof of ADMM specific instance of DR is tedious and error-prone. We thus give direct proof here instead. We first recall the DROT method:
X [Y, -pp[++
(13a)
Z = Pxx2X Yv))
(13b)
Y&+1 Ya + ZR+1 X
(13c)
Since Xx+1 is the projection of Yt pC onto RTxx" it holds that
(X Yv p,, X X AX R
(14)
Also, since X' is a closed affine subspace, we have
(Z++1 2X Yv, Z Z++1) =0 ZA E X
(15)
14

---

Published as a conference paper at ICLR 2022
Next, for X RT Z E X and U R mm we define the function:
V(X, 2,))
(U X).
Let Uk+1 (Yo (k+1)// //p, it holds that
( 1k+1) - 2, 1,1)- =s(X++1) + ( + (Uk+1, 7x+1 - Xx+1) - [(X) -(() - 1U 7 X)
1
=(C,XA+1) (C, - (Yx -x++1, 7x+1 - -7+X x x-xx+1)
p
By (14), we have (Ye- X X xx+1 > /p (C,X) - -CC,Xx+1> we thus arrive
1
V(X++1, Z4+1. Uk+1) - Z, Uk+1) v (Y:: -Xx+1, Zk+1 Z}
8
= (X ZR+1, Z4+1 Z) _ p
|| Zx+1 112 + (Xe+1 Zu+t, Xk+1 Z),
1
(16)
where the second step follows from (15). Now, for any Y Rmx,, we also have
v(Xx 1, Zx+1, Y - Z
V(X++1, 1 Uk+1)
P
1 (V Z, Z, Xk+1)= 1 (Yr -X411, Zkt1 -x+11
(Y - Yk. Zx+1 - -x+11) + (XK+1 - Z, -,,1+11 - -x+1). (17)
Therefore, by adding both sides of (16) and (17), we obtain
V( Z2+1, Y - 2, s (Y Yu, Zu+i XRt1)
1 ||x Zz41
p
(18)
Since Zx+1 -xa+1 = Yx+1 - Yk, it holds that
1 Y Yk, Zk+1 - Xk+1)
1
2p I - 2
||
r11++
+
|Y - Yoll.. (19)
2p
2p
Plugging (19) into (18) yields
1 Y - V(X, - - |rss r12
1
1
V(X++1, Ze+1,
|Ya-1
Y112
|| Yxll?
2p
2p
2p
(20)
which by dropping the last term on the right-hand side gives the first claim in the theorem.
Let Y be fixed-point of the mapping T and let X be solution of (1) defined from Y* in the
manner of Lemma 2.1. Invoking (20) with Z = X = X* E Rmx* n X and summing both sides of (20) over the iterations O, O....,K k 1 gives
(
( Y Y X*
1 1| -Y1²). (21)
Xc, (C,
2x
Xx
< Te 2p
11oo Y12
P
2p
where X₁ = Dt X,/k and 7/
-/
Z:/ k. Since
1
1
1
1
I 2|
Y -112=
|Yoll
Yall
+
(Y, Yoo
2p
2p
2p
2p
it follows that
(C,X))
+
1 7, Xx=
Yv Yo
1
Yo
222
1
||
2
_
1
(x*
Zx -x))
2ph
2ph
p
15

---

Published as a conference paper at ICLR 2022
Since Y is arbitrary in the preceding inequality, must have that
Yu
Yo
Z
Nx
0.
(22)
k
from which we deduce that
(c,
-
-
1
|oll
|
+ 1
+- - Zz - -x))
2ph
2pk:
1
-
|ool
+
|| 2xx - -Xall
2pk
P
1
Yool
||**|| || Yoll
(23)
2pk
pk
By Lemma B.I, we have
1 - - S |Y s..
11o - -Y'1,,
which implies that Finally, plugging (24) into (23) and (22) yields
\ - Yoll S 1 - Y || + 11% /2 - -Y
(24)
(c, (C,x*) S
1
||o|l 2
+
|| No y-11
2p
Yoll
2 Y Y
1 Xx|l
k
k
which concludes the proof.
B.2 PROOF THEOREM 2
We will employ a similar technique for proving linear convergence as Wang & Shroff (2017) used for ADMM on LPs. Although they, in contrast to our approach, handled the linear constraints via relaxation, show that the linear convergence also holds for our OT-customized splitting algorithm which relies on polytope projections. The main difference is that the dual variables are given implic- itly, rather than explicitly via the multipliers defined in the ADMM framework. As consequence, their proof needs to be adjusted to apply to our proposed algorithm.
To facilitate the analysis of the given splitting algorithm, we let 1, and f 1m and consider the following equivalent LP:
minimize X.ZER*XR
(C,X)
subject to
Ze
f =
(25)
X X21 0
Z = X.
The optimality conditions can be written on the form
Z' = p
(26a)
T f
(26b)
X*
(26c)
Z' = X*
(26d)
eeeT + fv S C
(26e)
(C,X') - p μ¹¹ T:²² 2²= = 0.
(26f)
This means that of optimal solutions, S* polyhedron
S' =2 -xx=xxx40 (X, () (2) 0, (2) <
16

---

Published as a conference paper at ICLR 2022
where Meg and Mn denote the equality and inequality constraints of (26), respectively. Li (1994) proposed a uniform bound of the distance between a point such a polyhedron in terms the con- straint violation given by:
M Meg 2 [Min 2)
dist(Z, - =g 0s.
(27)
For the (X, Z) )-varibbles generated by (9), (26a), (26b), and (26c) will always hold due the projec- tions carried out in the subproblems. As consequence, only (26d), (26e), and (26f) will contribute the bound (27). In particular:
Z -Xx (C, Mk 9 U/ [aRe + [vk - -C]+
dist((Xx; Zk, Mk, vk), 8*) As.
(28)
The following Lemma allows to express the right-hand side of (28) in terms of Yk+1 Lemma B.2. Let Xk. Zu and Ys be sequence generated by (9). It then holds that
Yk-
Zu+1 1++1-1++1 Yx+1 - Yx
MattE fu C' S p- (Y Y2) (C,X pk+1 9 Vk+1 = -pp- (Yx+1, Yu+1 Yk)
Proof. The first equality follows directly from the Y-update in (9). Let
P4 0u/P; Vs sk/p
(29)
the Y -updatee in (10) yields
Ye+1 = Yv + Zx+1 X6+1 Xat1 + P + +svktrr
(30)
In addition, the X-update
X = max(Ya -C,00
(31)
can be combined with the first equality in (30) to obtain
Y min(Y + Zk+1, Zu+1 pC).
equivalently
Yo+1 Zk+1 pC mmimminYYi PC,0).
(32)
By (30), the left-hand-sid in (32) can be written as
Yk+1 Za+1 pC Xx+1 - Zuti + 0p1+11e T pf PC
p(uuk11e" + fv -C C ' Vt)).
Substituting this in (32) gives
Pk+tel T fult1 C' CC Y2)) = pp min(
mnnin(Ys Yo pC, 0.
(33)
which gives the second relation in the lemma. Moreover, (31) and (33) implies that X and Matte + + p- -YY+( Yn)= C' cannot be nonzero simultaneously. As consequence,
(XK41, + futt - p- 2Y Yx) - C) =0
(C, Xx41) = (XEt1,MEHEE + SUI-I -Y+1 YY))
(Zu+1 - (Yu+1 Yx), +f - p - Yx))
17441,/4441e" + +funn) - (0-1 7k+1 + m T + futti, Ye+1 - Yx)
+ +p- ||Y++1 - 2
17

---

Published as a conference paper at ICLR 2022
By (30), p->)>> + Patie fuk+1 =p p Yk+1 - -x+11 p pp (2Y Ys), hence
Xx+1) = (00+4))=(Zu+tin) (Zuth + sub+22 - (2Y - Yn, Y641- - .aa) + p |Ye+1 - 1x
= + Sukt1) p- (Y++1: - - Yk)
(34)
Note that
(ZE4E + tr(Zk+1/441ee +
= -tr(e' (++116+11++++)))+)++++1)))1))1 + tr(v 1k+t
tr(p' (x+1) + +t((v(+19)
=p" Mk+1 + Mk+1-
(35)
Substituting (35) into (34), we get the third equality of the lemma:
(C, Xk-1) p p4+11 + Uk+1 =pp (Y++1, V2+1 - Ya)
By combining Lemma B.2 and (28), obtain the following result:
Lemma B.3. Let X2. Za and Ys be sequence generated by (9) and pax Vk be defined by (29). Then, there is y >> such that:
dist((Xk, Zu. Mk, vk), S \\r 11l
Proof. Let Y* E G", then Lemma B.1 asserts that
Y Y S Y Y s... S \| Y*I1,
which implies that |&|| < M for some positive constant M. In fact, {Yr. belongs to ball centered Y* with the radius Yo Y* From the bound (28) and Lemma B.2, we have
Z X2
dist((Xx, Zao Mkx S*) SAs.
(C, p ut q
[uge + +SV - -C+
0s (1 Xxll + +1(0,Xn) -p pp" 14k - V4l
+ [uke" + fuj - -C]+1l)
<0s. (1 - Y4-1ll +
+ p- p11 (Yx, Y, Yx-1) +p 1[Y2 Yx-1l+11)
50s. (1 + x- + p! \\x \ YE-
20s (1 + P + 1)) I - Y-1 ||.
By letting 7 = (1+01=+11) p (M + 1)) we obtain the desired result.
1
Recall that for DROT, prox, pf (-) (. pC) and prox pg (+) = Pxxx The following Lemma bridges the optimality conditions stated in (26) with G*
Lemma B.4. = + plateT + su+T), where X*, n²,, and u² satisfy (26), then G(Y+) 0.
Proof. We need to prove that
G(Y") = Papxa -8'-pC) Px 2P (2 (X* -00--Y' rY.)
equal to zero. Let us start with simplifying P PRIXx (Y* pC). By complementary slackness,
(X",
0,
and by dual feasibility
J f -C5
18

---

Published as a conference paper at ICLR 2022
Consequentially, have: quee" + fu² -C,X 50, XA E ER
which defines normal cone of R4x,, i.e. e²e + sO - -C NBmxx(X*) (X*). By using the definition of Y*, and that p 0, this implies that
- X pC E N2 Ngr.x (X*),
which corresponds to optimality condition of the projection
X PR (V* pC).
We thus have Px 22 Rmpx (Y* - - r+) Px(2X Y*). But by the definition of Y*, (Y+
0= - X plu' T + ff - -Y) -Y pp-e -pfv*
which means that X Px Y*) (following from the optimality condition of the projection). Substituting all these quantities into the formula for G(Y*), we obtain G(Y+) = 0.
Lemma B.5. There exists a constant e 2 such that
g") S cl|V2 Yk-1ll-
Proof. Let Y X* + puteT + pfv"T, where (X*, ,²",,")) v²* = Pss ((Xir Zko ukr Uk)) Ac- cording to Lemma B.4, we have that y* E G.. Then
dist(Yk: 9*) s ||4 - = ||Yr- X pllie + fftt)]l.
Since Yv = Xx + MALEE T + pf pfv it holds that
dist(Yr: 9*) s |X X* + p(ux - 4e + pf(ux - -1*)TI
Xx -X*1l + p|| n²1l + pll v*1l
$ /1+ p²||e|12 p² 1f X X+1 |1 - 1 2 + | - v=12
$ //||/+ + 11xx X+, Zx X :14 10, uk w+)1
/1c+ dist Zk. pkc vx), S*)
$ r1 + [|/ Yg-ill.
where the last inequality follows from Lemma B.3. The fact that 71 + p(|le|l + (lflD) 2 1
(36)
follows from (12a) in Corollary B.
L
B.3 PROOF OF LINEAR CONVERGENCE
By combining Lemma B.5 with (12b) in Corollary B. we obtain:
2 2
(Yist(Y)+1,
(dist(Yk, g*))²
or equivalently
dist(Ve+1, < VI+ dist((ys (k, g*).
Letting r = c/v11 c² and iterating the inequality times gives dist(Yk, g*) dist(Yo. (ist Let V Pg. (Yk). Since prox,f (Y*) E X and X prox pf (Yx), we have
dist dist(Xx+1, ||pox,, (Yr) prox,, (Y*) | Y dist(Yk. 9+) S dist(Yo. )* where the second inequality follows from the non con-expansiveness of prox, (.) This completes the proof of Theorem 2.
19

---

Published as a conference paper at ICLR 2022
C DR AND ITS CONNECTION TO ADMM
The Rachford updates given (5) can equivalently formulated as:
Ik+1 prox pf (1k)
#k+1 prox (2TR+1 + 11b)
Mk+1 y/ + 21+1 3k++1
+
2k+1 proXpg (2IR+1 + yk)
uk+1
yk + 2k+1 2x+1
Tk42 prox ps (yk+1)
+
2k41 proxpg (2I1+1 1 + 1kb)
Tk42 prox, (y1 + 2b+1 22+1)
1k+1 Y) + 2 2x+1.
If we let wk+1 (yxx 2k+1), get.
2k+1 = proXpg (2k+1 + pu6++1
Tk+2 = prox,s (#k+1 - p(0k+1)
Wk+2 wk+1 + (2k+1 x1+2))
By re-indexing the x and the wo variables, we get:
1 prox (It + pww)
Tk+1 prox (#k+1 pwk)
-
Wi+1 Wk (EK+1 T1+1).
This is exactly ADMM applied to the composite convex problem on the form (4).
D TENCHEL-ROCKAFELLLA DUALITY OT
Recall that
S(X) (C, X) I==xx= and g(X) = (X)
=
Recall also that the conjugate of the indicator function le of closed and convex set C is the support function de of the same set (Beck, 2017, Chapter 4). For a non-empty affine set X' = (Y A(Y) b}. its support function can be computed (U = (U, Xo) + IranA' (U),
where No is point in X' (Beck, 2017, Example 2.30). By letting Xo pqj E X, it follows that g gu( = (U,pq (U))
Since A gu + Rmxx; (v,z) ye + Tx any matrix U E ranA* must have the form
ue" T + fv for some μ E Im and v E R The function g" (U)) can thus be evaluated as:
g0( = p qT Tv.
We also have
f²(-)) mrxx((-U,X) - (C,,) - I (X)} x
max((-U CC,X) - 1m (X))
Ogmnn (-U (-m -C)
- C
={ t00
otherwise.
20

---

Published as a conference paper at ICLR 2022
The dual of problem 4 in Lemma 2.1 thus becomes
maximize PERm ME
-p" μ q- v subject to _ee- - fv - 0,
which is identical to (3) upon replacing μ by - and v by -v.
21


Reviews:
[{"review_id": "qNzOaG8RTGy", "reviewer": "Reviewer_6N6U", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "This paper report a new efficient numerical scheme for the fundamental discrete OT problem. The main technical contribution is a new splitting scheme by a recent result on efficient range space projection, which is very interesting. The authors also parallelize the algorithm for GPU computing and made the procedure numerical stable and efficient. The theoretical results are somehow standard as DR splitting is a special case of ADMM, of which the standard convergence behavior is well understood now. The linear convergence is mainly due to the KKT mapping being polyhedral and thus metric subregular. Existing results could directly give the linear convergence.\n\nMy main questions are as follows:\n\n* The experiments are mostly conducted compared with the Sinkhorn method. It would be interesting to see the results of other splitting schemes for discrete OT, e.g., the iterative Bregman projection.\n\n* It would be good to also provide empirical objective function value convergence curves. As the global linear convergence (due to global Hoffman bound) is proved, it would be interesting to see the empirical performance for that.\n\n* Could the 1/k convergence be made non-ergodic?\n", "summary_of_the_paper": "This paper introduces a new splitting scheme for the original discrete OT problem. Unlike the existing splitting techniques, their solution requires an efficient implementation of rectangular matrix range space projection. The authors provide (somehow standard) 1/T convergence complexity and global linear convergence complexity. They report numerical results compared with the Sinkhorn method.", "summary_of_the_review": "This paper introduces a new efficient DR scheme for discrete OT. The problem is important as it's widely used in ML models. The new splitting Eq.(7) is interesting, and the authors have spent an extensive effort to make the implementation efficient and robust.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "kyo591tLoZK", "reviewer": "Reviewer_szP9", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "While alternate minimization type approaches for OT have been extensively studied, I find the performance of about 1 second for 20000-by-20000 instances highly impressive. This requires a fairly high degree of parallelism, and is quite difficulty to obtain without significantly adapting to hardware. Some of the subtleties, such as numerical instability of sinkhorn2, is also something that I have observed before, so I consider the experiments quite thorough.\n\nThe algorithms are theoretically well founded, and have convergence guarantees typical of this type of methods. They also have the advantage of working directly with regularized OT instances, which to my understanding are critical for the downstream applications using OT. My main worry about these theoretical guarantees is that they are quite difficult to compare vs. unconditional bounds such as those obtained via combinatorial optimization packages, or the recent (purely theoretical) n^2 type methods (https://arxiv.org/pdf/2101.05719.pdf).", "summary_of_the_paper": "This paper studies a splitting based gradient descent method for computing optimal transport objectives. It's based on splitting the row/column sum constraints into two objectives, and using a variant of the DR splitting algorithm. The paper rigorously proves a convergence guarantee of 1/k times the initial distance after k steps. Then it explores ways of setting up its algorithm in ways friendly to parallelism / memory hierarchy empirically: they obtain a very good wall-clock performance of about 1 second on a 20,000-by-20,000 optimal transport instance, and also examine the convergence of these algorithms.", "summary_of_the_review": "Overall, I believe this paper has useful ideas for improving the performance of OT methods, especially those currently in use, but it falls a bit short of proposing something that significantly advances the state-of-art for OT.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "_-cQFtPWVXG", "reviewer": "Reviewer_Jdt4", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The authors address an important problem and come up with an interesting algorithm. They also provide a rigorous theoretical convergence analysis as well as many insights on their CUDA implementation. The performance of the proposed method is competitive with Sinkhorn.\n\nUnfortunately this paper also has some major limitations. The biggest one lies in its rather modest set of experiments. Some fundamentally important experiments are missing to support the claims of the paper.\n\n1. Implementation is presented as a core contribution of this paper: it is mentioned in the title, and in addition, it is devoted a substantial amount of space for presentation (almost 3 pages, a part of Section 3.1 and especially the whole Section 3.3). Therefore, it would be absolutely necessary to show that the presented materials are important by actually comparing the provided implementation (with highly-optimized CUDA kernels) with a naive implementation, e.g., using GPU linear algebra libraries such as PyTorch or TensorFlow. The proposed method (Algorithm 1) can be easily implemented using these libraries in just a few lines.\n\n2. The color transfer experiment is weird. Why aren't the results of Sinkhorn presented? Moreover, if sparsity is considered as an important objective, then the authors should also compare with Blondel et al. (2018). I would suggest to consider a more diverse set of examples to showcase the performance of DR in comparison with Sinkhorn (see also the next point).\n\n3. Sinkhorn has been highly successful not only because of its efficiency, but also because its updates are differentiable, which makes it suitable for learning with stochastic gradient descent. While the updates of Algorithm 1 are also differentiable almost everywhere, it is unclear how this algorithm performs in learning, compared to Sinkhorn. I would encourage the authors to add some learning experiments to make the paper even more solid.\n\n\nFinally, this paper also has some presentation issues.\n\n1. Major issues:\n\n- I wonder why the authors decided to denote 1_n as \"e\" and 1_m as \"f\". The notation \"f\" is used at multiple locations with different meanings (while the \"e\" can be confused with the Euler's number, though this is less critical than the \"f\"). It would be laborious to change because this is used everywhere, so I would suggest (as an acceptable fix) to replace the other instances of \"f\" with \"h\", e.g., \"h(x) + g(x)\" instead of \"f(x) + g(x)\".\n\n- In Figure 1 and 2, the descriptions of the axes should be clear from the captions (instead of being hidden in the text). Furthermore, the text in the plots (such as axis labels, markers, legends) are unreadable (too small).\n\n- Figure 3b shows that Sinkhorn is actually faster to reach a given accuracy (even though it is less numerically stable). This would deserve some comments (in a fair comparison of two methods, we should present both strengths and weaknesses of both methods).\n\n2. Minor issues:\n\n- Page 2, last paragraph of Section 1.1: \"Liang et al. (2017), derived\" --> \"Liang et al. (2017) derived\"\n\n- Page 2, line -5: \"methods(Rockafellar, 1976)\" --> \"methods (Rockafellar, 1976)\"\n\n- Section 2 2nd paragraph: \n\t- \"computed as Bauschke et al. (2021):\"  --> \"computed as (Bauschke et al., 2021):\"\n\t- \"linear program on the form\" --> \"linear program of the form\"\n\n- Page 3 last paragraph: \"problems on the form\" --> \"problems of the form\"\n\n- Section 3, 2nd paragraph: \"and fact that\" --> \"and the fact that\"\n\n- The notation for indicator function in Eq. (7) should be defined.\n\n- Page 5 first paragraph: \"matrix-vector multiples\" --> \"matrix-vector multiplications\"\n\n- Page 8 last paragraph: \"A algorithm\" --> \"An algorithm\"\n\n---\n**Update on November 10th:**\n\nI would like to have some additional comments/questions. Hopefully the authors could address them as well in their rebuttal:\n\n1. There seems to be a minor (likely fixable) **flaw** in Theorem 1 and the discussion thereafter. The inequality \n\\begin{equation}\n\\langle C,\\bar{X}\\_k \\rangle - \\langle C,X^* \\rangle \\le \\mathcal{O}(\\|Y\\_0-Y^*\\|^2/k)\\qquad (\\star)\n\\end{equation}\n\nis **not sufficient** to yield the convergence of the objective (let alone its rate). Note that **$\\bar{X}_k$ might not be feasible** and thus $\\langle C,\\bar{X}\\_k \\rangle - \\langle C,X^* \\rangle$ might be negative. In order to obtain the convergence of the objective, one would need a lower bound as well:\n\\begin{equation}\n\\alpha_k \\le \\langle C,\\bar{X}\\_k \\rangle - \\langle C,X^* \\rangle \\le \\mathcal{O}(\\|Y\\_0-Y^*\\|^2/k),\\qquad (\\star\\star)\n\\end{equation}\nwhere $\\alpha_k\\to 0$. Furthermore, in order to obtain the claimed $\\mathcal{O}(1/k)$ rate of convergence, one would also need to show that $(\\alpha_k)$ converges to $0$ at the same rate.\n\n2. The authors mentioned that they use the Sinkhorn implementation from POT. It is important to provide enough details so that the interested reader could reproduce the results. For example, which function of POT did the authors use and with which arguments? (ot.sinkhorn has quite several options for its arguments.)", "summary_of_the_paper": "This paper proposes a novel method for solving the discrete optimal transport problem based on the celebrated Douglas-Rachford splitting (DR) method. The core idea is to split the original transportation map into two variables X and Z, where X is simply non-negative and Z is real-valued and obeys the given row and column sums constraints. The Z-update step in DR is then non-trivial, consisting in projecting a given matrix onto the set of all real-valued matrices with such row/column constraints, which has a closed-form solution due to Bauschke et al. (2021). The authors present a detailed discussion on how to implement the resulting algorithms efficiently on GPUs, and provide a theoretical convergence analysis that recovers the known O(1/k) rate of convergence for DR method. Experiments on a synthetic problem show that the proposed method (using the authors' CUDA implementation) is more numerically stable than the standard Sinkhorn method, while being comparably fast.", "summary_of_the_review": "The proposed algorithm and analysis/implementation are interesting but the experiments are not convincing. The paper also has some issues with the presentation (but these can be easily fixed).", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "Ph9gN5dmTC0", "reviewer": "Reviewer_qeU8", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "- It's a novel idea to rewrite the original OT problem into the standard form for DR-splitting and exploit the DR-splitting algorithm to solve it. The DROT method performs well in aspects of both speed and accuracy, thus being a strong competitor of the popular Sinkhorn method. \n\n- The proposed algorithm can save memory efficiently for it can be executed with only a single matrix variable.\n\n- The organization of this paper is satisfactory, and the flow of algorithm derivation is very clear.\n\n- Experiments on much larger-scale datasets are needed to verify the effectiveness of the proposed OT solver.\n\n- Log-domain stabilization of Sinkhorn proposed by [1] is also a popular strategy to deal with numerical issues introduced by entropy regularization, while there is a lack of comparison with it.\n\n[1] Schmitzer, Bernhard. \"Stabilized sparse scaling algorithms for entropy regularized transport problems.\" SIAM Journal on Scientific Computing 41.3 (2019): A1443-A1481.\n\nThere are some typos and nitpicks:\n\n- In line 1 of page 2, $\\|v_{k} \\odot\\left(K^{\\top} u_{k}\\right)-q\\|$ should be $\\|v_{k-1} \\odot\\left(K^{\\top} u_{k}\\right)-q\\|$. Note that $\\|v_{k} \\odot\\left(K^{\\top} u_{k}\\right)-q\\|$ always equals to 0 due to the Sinkhorn iteration equation.\n\n- In figure 3, it’s better not to use colors to distinguish various methods, as this is not friendly for color-blind people; the text font of labels and axes is too small to read.\n\n\n\n", "summary_of_the_paper": "In order to solve large-scale optimal transport problems, the authors propose a Douglas-Rachford (DR) splitting method that solves the original OT without entropy regularization. \nThe proposed DROT algorithm has the same cost per iteration as the well-known Sinkhorn method but avoids its numerical issues.\nFurthermore, the proposed method possesses the strong convergence guarantees for its linear convergence rate; and it can be implemented efficiently using parallel processing on GPUs.\nExperiments verify its efficiency and accuracy compared to Sinkhorn.\n", "summary_of_the_review": "The proposed method is interesting, while some points listed above need to be clarified.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper proposes a new algorithm to solve the discrete optimal transport problem, consisting of showing how the well-known Douglas Rachford algorithm can be efficiently applied, and also providing a convergence rate. Secondly, the paper gives an efficient implementation suitable for GPUs.

The paper would be a bit weak on just the DR contribution alone, so the efficient implementation, and experiments, are significant. However, the reviewers were not consistently happy with the experiments.

Two big issues raised by reviewers were wanting more experiments (having comparison experiments with real-world data), and wanting to compare with variants of the baseline algorithm (Sinkhorn), such as the log-transformed version which is more stable.  Looking at the revision, I think the reviewers have done a good job adding experiments.  Overall, for a paper with strong theoretical components, I think the computational aspects are strong.

As for (not) comparing with the log-transformed Sinkhorn and other more robust methods, the authors argue that this implementation would be slower.  I agree with reviewers that it would be nice to have these comparisons, but the authors' argument is plausible and I don't find it grounds to reject the paper.

Overall, it seems there is some evidence that this is a worthwhile method, and there are no theoretical concerns other than presentation issues. Thus I think it would be a benefit to the community to accept this paper.
Decision: Accept (Poster)