A FAST AND ACCURATE SPLITTING METHOD FOR OPTIMAL TRANSPORT : ANALYSIS AND IMPLEMENTATION
Jacob Lindb¨ack
Mikael Johansson †
A BSTRACT
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 efﬁciently, in parallel.
The proposed method enjoys an iteration complexity
O
(1 /(cid:15) )
compared to the best-known O (1 /(cid:15) 2 ) of the Sinkhorn method.
In addition, we establish a linear convergence rate for our formulation of the OT problem.
We detail an efﬁcient 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.
I NTRODUCTION
We study the discrete optimal transport (OT) problem: minimize
[Equation]
subject to X 1 n = p and X (cid:62) 1 m
[Equation]
q [Equation] [Equation] where X ∈ R m × n + is the transport plan, C ∈ R m × n + is the cost matrix and p ∈ R m + and q ∈ R n + 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 ﬂexible way to compare probability measures.
It has quickly become a 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 et al., 2019); we refer to Peyr ´ e & 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-linear time complexity in term of the problem size.
Such methods are also very challenging to parallelize on modern processing hardware.
Therefore, there has been substantial research in developing new efﬁcient methods for OT.
This paper advances the state of the art in this direction.
1. R ELATED 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 a regularized version of the form (cid:104) C, X (cid:105) − ηH ( X ) .
Here,
H ( X ) = − (cid:80) ij X ij log( X ij ) is an entropy function and η > 0 is a regularization parameter.
The
Sinkhorn method deﬁnes the quantity K = exp( − C/η ) and repeats the following steps u k =
p/ ( Kv k − 1 ) and v k =
q/ ( K (cid:62) u k ) , until (cid:107) u k (cid:12) ( Kv k ) − p (cid:107) + (cid:107) v k − 1 (cid:12) ( K (cid:62) u k ) − q (cid:107) becomes small, then returns diag( u k ) K diag( v k ) .
The division ( / ) and multiplication ( (cid:12) ) operators between two vectors are to be understood entrywise.
Each SK iteration is built from matrix-vector multiplies and element-wise arithmetic operations, 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 to ﬁnd 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 η that is small enough to provide a meaningful approximation, and large enough to avoid numerical issues.
In addition, the entropy term enforces a dense solution, which can be undesirable when the optimal transportation plan itself is 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;
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 a 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 a relatively low per-iteration cost and the ability to exploit sparsity in the problem data, operator splitting methods such as Douglas-Rachford splitting (Douglas & Rachford, 1956) and ADMM (Gabay & Mercier, 1976) have gained widespread popularity in large-scale optimization.
Such algorithms can quickly produce solutions of moderate accuracy and are the engine of several successful ﬁrst-order solvers (O’donoghue 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 splitting-based 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 bottleneck is not ﬂoating-point computations, but rather time-consuming memory operations on large two-dimensional arrays.
Even an innocent update like X ← X − C is more expensive than the two matrix-vector multiplies in SK.
To design a 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 prohibitively 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 e.g. Giselsson & Boyd (2016)).
Recently, it has been shown that algorithms based on DR/ADMM often enjoy similar convergence properties also in the absence of strong convexity.
For example, Liang et 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) established linear convergence for an ADMM-based algorithm on general LPs.
Yet, these frameworks quickly become intractable on large OT problems, due to the many memory operations required.
1. C ONTRIBUTIONS
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.
Speciﬁcally, 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 simpliﬁcations to eliminate variables so that the ﬁnal algorithm can be executed with only one matrix variable, while maintaining the same degree of parallelization.
Our method implicitly maintains a primal-dual pair, which facilitates a simple evaluation of stopping criteria.
• We derive an iteration complexity O (1 /(cid:15) ) for our method.
This is a signiﬁcant improvement over the best-known estimate O (1 /(cid:15) 2 ) for the Sinkhorn method (cf.
Lin et al.
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 an efﬁcient GPU implementation that fuses many immediate steps into one and performs several on-the-ﬂy reductions between a read and write of the matrix variable.
We also show how a primal-dual stopping criterion can be evaluated at no extra cost.
The implementation is
available as open source and gives practitioners a fast and robust OT solver also for applications where regularized OT is not suitable.
As a 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 ﬁnite time, even before reaching a solution Iutzeler & 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 is readily extended to other splitting schemes (e.g. Chambolle & Pock (2011)).
B ACKGROUND
Notation For any x, y ∈ R n , (cid:104) x, y (cid:105) is the Euclidean inner product of x and y , and (cid:107)·(cid:107) 2 denotes the (cid:96) 2 -norm.
For matrices X, Y ∈ R m × n , (cid:104) X, Y (cid:105) =
tr( X (cid:62) Y ) denotes their inner product and (cid:107)·(cid:107) F = (cid:112) (cid:104)· , ·(cid:105) is the the induced Frobenius norm.
We use (cid:107)·(cid:107) to indicate either (cid:107)·(cid:107) 2 or (cid:107)·(cid:107) F .
For a closed and convex set X , the distance and the projection map are given by dist( x, X ) =
min
z ∈X (cid:107)
z − x (cid:107) and P X ( x ) = argmin
z ∈X (cid:107)
z − x (cid:107) , respectively.
Further, we denote the proximal operator of a closed convex function f by prox ρf ( x ) = argmin
z f ( z ) + 12 ρ (cid:107) z − x (cid:107) 2 , ρ > 0 . The Euclidean projection of x ∈ R n onto the nonnegative orthant is denoted by [ x ] + = max( x, 0) , and ∆ n = { x ∈ R n + : (cid:80) ni
=1 x i = 1 } is the ( n − 1) dimensional probability simplex.
OT and optimality conditions Let e = 1 n and f = 1 m , and consider the linear mapping
[Equation]
and its adjoint A ∗ : R m + n → R m × n : ( y, x ) (cid:55)→ ye (cid:62) + fx (cid:62) .
The projection onto the range of A is P ran A (( y, x )) =
( y, x ) − α
( f, − e ) , where α = ( f (cid:62) y − e (cid:62) x ) /
( m + n ) (Bauschke et al., 2021).
With b = ( p, q ) ∈ R m + n , Problem (1) can be written as a linear program on the form
[Equation]
Let ( µ, ν ) ∈
R m + n be the dual variable associated the afﬁne constraint in (2).
Then, using the deﬁnition of A ∗ , the dual problem of (2), or equivalently of (1), reads
[Equation]
OT is a bounded program and always admits a 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 ( µ, ν ) are primal and dual optimal if and only if: (i)
Xe
= p, X (cid:62) f
= q, X ≥ 0 , (ii) µe (cid:62) + fν (cid:62) − C ≤ 0 , (iii) (cid:104) X, C − µe (cid:62) − f (cid:62) [Equation] (cid:105) =
[Equation]
These conditions mean that: (i) X is primal feasible; (ii) µ, ν are dual feasible, and (iii) the duality gap is zero (cid:104) C, X (cid:105)
= p (cid:62) µ + q (cid:62) ν , or equivalently, complementary slackness holds.
Thus, solving the OT problem amounts to (approximately) ﬁnding such a primal-dual pair.
To this end, we rely on the celebrated Douglas-Rachford splitting method (Lions & Mercier, 1979; Eckstein & Bertsekas, 1992; Fukushima, 1996; Bauschke & Combettes, 2017).
Douglas-Rachford splitting Consider composite convex optimization problems on the form
[Equation]
where f and g are proper closed and convex functions.
To solve Problem (4), the DR splitting algorithm starts from y 0 ∈ R n and repeats the following steps:
[Equation]
where ρ > 0 is a penalty parameter; prox ρf ( · ) and prox ρg ( · ) are the proximal operator of f and g , respectively.
Procedure (5) can be viewed as a ﬁxed-point iteration
y k +1 =
T ( y k ) for the mapping
[Equation]
The DR iterations (5) can also be derived from the ADMM method applied to an 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 a ﬁrmly-nonexpansive mapping, from which they obtained convergence of y k .
Moreover, the sequence x k is guaranteed to converge to a minimizer of f + g (assuming a minimizer exists) for any choice of ρ > 0 .
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–Rockafellar dual deﬁned as
[Equation]
Let P (cid:63) and D (cid:63) be the sets of solutions to the primal (4) and dual (7) problems, respectively.
Let x k , y k , and z k be generated by procedure (5) and let u k :=
( y k − 1 − x k ) /ρ .
Then, there exists y (cid:63) ∈ R n such that y k → y (cid:63) .
Setting x (cid:63) =
prox ρf
( y (cid:63) ) and u (cid:63) =
( y (cid:63) − x (cid:63) )
/ρ , then it holds that (i) x (cid:63) ∈ P (cid:63) and u (cid:63) ∈ D (cid:63) ; (ii) x k − z k
→ 0 , x k → x (cid:63) and z k → x (cid:63) ; (iii) u k → u (cid:63) .
D OUGLAS -R ACHFORD SPLITTING FOR OPTIMAL TRANSPORT
To efﬁciently apply DR to OT, we need to specify the functions f and g as well as how to evaluate their proximal operations.
We begin by introducing a result for computing the projection onto the set of real-valued matrices with prescribed row and column sums (Romero, 1990; Bauschke et al., 2021).
Lemma 3.1.
Let e = 1 n and f = 1 m .
Let p
∈
∆ m and q ∈ ∆ n .
Then, the set X deﬁned by:
[Equation]
is non-empty, and for every given X ∈ R m × n , we have
[Equation]
where γ = f (cid:62) ( Xe − p ) / ( m + n ) =
e (cid:62) (cid:0) X (cid:62) f − q (cid:1) /
( m + n ) .
The lemma follows immediately from Bauschke et al. (2021, Corollary 3.4) and the fact that ( p, q ) =
P ran A (( p, q )) .
It implies that P X ( 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. A LGORITHM DERIVATION
Our algorithm is based on re-writing (2) on the standard form for DR-splitting (4) using a carefully selected f and g that ensures a rapid convergence of the iterates and an efﬁcient execution of the iterations.
In particular, we propose to select f and g as follows:
[Equation]
By Lemma 3.1, we readily have the explicit formula for the proximal operator of g , namely, prox ρg ( · ) =
P X ( · ) .
The proximal operator of f can also be evaluated explicitly as:
[Equation]
The Douglas-Rachford splitting applied to this formulation of the OT problem then reads:
[Equation]
Despite their apparent simplicity, the updates in (9) are inefﬁcient to execute in practice due to the many memory operations needed to operate on the large arrays X k , Y k , Z k and C .
To reduce the memory access, we will now perform several simpliﬁcations to eliminate variables from (9).
The resulting algorithm can be executed with only one matrix variable while maintaining the same degree of parallelization.
We ﬁrst note that the linearity of P X ( · ) allows us to eliminate Z , yielding the Y -update
[Equation]
Input: OT( C , p , q ), initial point X 0 , penalty parameter ρ 1: φ 0 = 0 , ϕ 0 =2: a 0 = X 0 e − p , b 0 =
X
{+(cid:62) 
+}
f − q , α 0 = f (cid:62) a 0 / ( m + n ) 3: for k = 0 , 1 , 2 , . . .
do 4:
X k +1 = (cid:2)
X k
+ φ k e (cid:62) + fϕ (cid:62) k − ρC (cid:3) + 5:
r k +1 =
X
k +1 e − p , s k +1 = X (cid:62) k +1 f − q , β k
{++
+}
=
{+f (cid:62) r k +1 / ( m + n ) 6: φ k +
+}
=
{+( a k − 2 r k +
+}
+ (2 β k +1 − α k ) f )
/n 7: ϕ k
+1 = ( {+b k − 2 s k +
+} + (2 β k +1 − α k ) e ) /m 8: a k +1 =
a k − r k +1 , b k +1 =
b k − s k
{++1 , α k +
+}
=
α k − β k + 9: end for Output: X K where γ k = f (cid:62) {+(2 X k +
 e+}
− Y k
e − p ) / ( m + n ) =
e (cid:62) (cid:0) 2 X (cid:62) k +1 f − Y (cid:62)
k f − q (cid:1) /
( m + n ) .
We also deﬁne the following quantities that capture how Y k affects the update of Y k + a k =
Y k e − p, b k = Y (cid:62)
k f − q, α k = f (cid:62) a k / ( m + n ) =
e (cid:62) b k / ( m + n ) .
Similarly, for X k , we let:
[Equation]
Recall that the pair ( r k , s k ) represents the primal residual at X k .
Now, the preceding update can be written compactly as
[Equation]
where
[Equation]
We can see that the Y -update can be implemented using 4 matrix-vector multiples (for computing a k , b k , r k +1 , s k +1 ), followed by 2 rank-one updates.
As a rank-one update requires a read from an input matrix and a write to an output matrix, it is typically twice as costly as a matrix-vector multiply (which only writes the output to a vector).
Thus, it would involve 8 memory operations of large arrays, which is still signiﬁcant.
Next, we show that the Y -update can be removed too.
Noticing that updating Y k +1 from Y k and
X k +1 does not need the full matrix Y k , but only the ability to compute a k and b k .
This allows us to use a single physical memory array to represent both the sequences X k and Y k .
Suppose that we overwrite the matrix X k
+
as:
[Equation]
then after the two rank-one updates, the X -array holds the value of Y k
+1 .
We can access the actual
X -value again in the next update, which now reads: X k +2 ← (cid:2)
X
k +1 − ρC (cid:3)
+ .
It thus remains to show that a k and b k can be computed efﬁciently.
By multiplying both sides of (10) by e and subtracting the result from p , we obtain
[Equation]
Since e (cid:62) e = n and ( b k − 2 s k +1 ) (cid:62) e = ( m + n )( α k − 2 β k +1 ) , it holds that ( ϕ (cid:62)
k
+1 e ) f
= ( α k − 2 β k +1 ) f .
We also have φ k +1 e (cid:62)
e = a k − 2 r k +1 + (2 β k
+
− α k ) f .
Therefore, we end up with an extremely simple recursive form for updating a k :
[Equation]
Similarly, we have b k +
=
b k − s k +
and α k +
=
α k − β k +1 .
In summary, the DROT method can be implemented with a single matrix variable as summarized in Algorithm 1.
Stopping criteria It is interesting to note that while DROT directly solves to the primal problem (1), it maintains a pair of vectors that implicitly plays the role of the dual variables µ and ν in the dual problem (3).
To get a feel for this, we note that the optimality conditions in Propositionare equivalent to the existence of a pair X (cid:63) and ( µ (cid:63) , ν (cid:63) ) such that
[Equation]
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(ii)) with the convex cone
C = R m × n + , y = X (cid:63) ∈ C and z = µ (cid:63) e (cid:62) + fν (cid:63) (cid:62) − C ∈ C ◦ .
Now, compared to Step 4 in DROT, the second condition above suggests that φ k /ρ and ϕ k /ρ are such implicit variables.
To see why this is indeed the case, let U k
= ( Y k − 1 − X k ) /ρ .
Then it is easy to verify that
[Equation]
By Lemma 2.1, we have Z k − X k[-+
-] → 0 and U k[-+
-]
→ U (cid:63) ∈ R m × n , by which it follows that
[Equation]
Finally, by evaluating the conjugate functions f ∗ and g ∗ in Lemma 2.1, it can be shown that U (cid:63) must have the form µ (cid:63) e (cid:62) + fν (cid:63) (cid:62) , where µ (cid:63) ∈ R m and ν (cid:63) ∈ R n satisfying µ (cid:63) e (cid:62) + fν (cid:63) (cid:62) ≤ C ; see Appendix D for details.
With such primal-dual pairs at our disposal, we can now explicitly evaluate their distance to the set of solutions laid out in Proposition 1 by considering: r primal = ( (cid:107) r k
(cid:107)
+ (cid:107) s k (cid:107) 2 ) 1 / 2 , r dual
= (cid:107) [( φ k /ρ )
e (cid:62) + f ( ϕ k /ρ )
(cid:62) − C ]
+ (cid:107) , and gap = (cid:12)(cid:12) (cid:104) C,
X k (cid:105) − ( p (cid:62) φ k + ϕ (cid:62) k q ) /ρ
(cid:12)(cid:12) .
As problem (1) is feasible and bounded, Lemma 2.1 and strong duality guarantees that all the three terms will converge to zero.
Thus, we terminate DROT when these become smaller than some user-speciﬁed tolerances.
3. C ONVERGENCE RATES
In this section, we state the main convergence results of the paper, namely a sublinear and a linear rate of convergence of the given splitting algorithm.
In order to establish the sublinear rate, we need the following function:
[Equation]
which is deﬁned for X ∈ R m × n + , Z ∈ X and U ∈ R m × n .
We can now state the ﬁrst result.
Theorem 1.
Let X k , Y k , Z k be the generated by procedure (9) .
Then, for any X ∈ R m × n + , Z ∈ X , and Y ∈ R m × n , we have V ( X
[Equation]
, Z
[Equation]
+1 , [Equation] (2 ρ ) . Furthermore, let Y (cid:63) be a ﬁxed-point of T in (6) and let X (cid:63) be a solution of (1) deﬁned from Y (cid:63) in the manner of Lemma 2.1.
Then, it holds that
[Equation]
where ¯ X k = (cid:80) ki
=
X i /k and ¯ Z k = (cid:80) ki
=1 Z i /k .
The theorem implies that one can compute a solution satisfying (cid:104) C, X (cid:105) − (cid:104) C, X (cid:63) (cid:105) ≤ (cid:15) in O (1 /(cid:15) )
iterations.
This is a signiﬁcant improvement over the best-known iteration complexity O (1 /(cid:15) 2 ) of the Sinkhorn method (cf.
Lin et al.
Note that the linearity of (cid:104) C, ·(cid:105) allows to update the scalar value (cid:10) C, ¯ X k (cid:11) recursively without ever needing to form the ergodic sequence ¯ X k .
Yet, in terms of rate, this result is still conservative, as the next theorem shows.
Theorem 2.
Let X k and Y k be generated by (9) .
Let G (cid:63) be the set of ﬁxed points of T in (6) and let X (cid:63) be the set of primal solutions to (1) .
Then, { Y k } is bounded, (cid:107) Y k (cid:107) ≤ M for all k , and dist( Y k , G (cid:63) ) ≤ dist( Y 0 , G (cid:63) ) × r k and dist( X k , X (cid:63) ) ≤ dist
( Y 0 , G (cid:63) )
× r k − 1 , where r = c/ √ c 2 + 1 < 1 , c = γ (1 + ρ ( (cid:107) e (cid:107) + (cid:107) f (cid:107) )) ≥ 1 , and γ = θ S (cid:63) (1 + ρ −
( M + 1)) .
Here, θ S (cid:63) > 0 is a problem-dependent constant characterized by the primal-dual solution sets only. [Figure]
This means that an (cid:15) -optimal solution can be computed in O log 1 /(cid:15) ) iterations.
However, it is, in general, difﬁcult 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 ﬁrst iterations.
In either case, DROT always satisﬁes the better of the two bounds at each k .
3. I MPLEMENTATION
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, a large number of threads are generated to execute its statements.
These threads are organized into a two-level hierarchy.
A grid contains multiple blocks and a block contains multiple threads.
Each block is scheduled to one of the streaming multiprocessors (SMs) on the GPU concurrently or sequentially, depending on available hardware.
While all threads in a thread block run logically in parallel, not all of them can run physically at the same time.
As a result, different threads in a thread block may progress at a different pace.
Once a 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 at the same time .
Each thread has its own instruction address counter and register state, and carries out the current instruction on its own data.
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 in a kernel is generally stored in a register.
Shared 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 a 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 on a GPU card is global memory , physically separated from the compute chip (“offchip”).
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 a high-performance kernel.
When all threads of a warp execute a 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 ﬂoat32 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 × n arrays.
In particular, this includes one read from and write to X , and one read from C in every other iteration.
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 efﬁcient kernel that: (i) updates X k to
X k +1 , (ii) computes u k +1 :
=
X k +1 e
and v k[-+
-]
+1 :
= X (cid:62) k +1 f , (iii) evaluates (cid:104) C, X k +1 (cid:105) in the duality gap expression.
The kernel fuses many immediate steps into one and performs several on-the-ﬂy reductions while updating X k , thereby minimizing global memory access.
Our kernel is designed so that each thread block handles a submatrix x of size bs × ( ws · bs ) , except for the corner blocks which will have fewer rows and/or columns.
Since all threads in a block need the same values of ϕ , it is best to read these into shared memory once per block and then let threads access them from there.
We, therefore, divide ϕ into chunks of the block size and set up a 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 ϕ and φ .
In the j -th step, the working block loads column j of x to the chip in a coalesced fashion.
Each thread i ∈ { 0 , 1 , . .
, bs − 1 } uses the loaded x ij to compute and store x
+ ij in a register:
[Equation]
where c is the corresponding submatrix of C .
As sum reduction is order-independent, it can be done locally at various levels, and local results can then be combined to produce the ﬁnal value.
We, therefore, reuse x + ij and perform several such partial reductions.
First, to compute the local value for u k +1 , at column j , thread i simply adds x + ij to a running sum kept in a register.
The reduction leading to (cid:104) C, X k +1 (cid:105) can be done in the same way.
The vertical reduction to compute v k +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 efﬁcient CUDA’s built-in functions for collective communication at warp-level. {+
+}
When done, the ﬁrst thread in a warp holds the total reduced value and simply adds it atomically to a proper coordinate of v k +1 .
Finally, all threads write x + ij to the global memory.
The process is repeated by moving to the next column.
When the block ﬁnishes processing the last column of the submatrix x , each thread adds its running sum atomically to a proper coordinate of u k +1 .
They then perform one vertical (warp-level) reduction to collect their private sums to obtain the partial cost value (cid:104) c, x (cid:105) .
When all the submatrices have been processed, we are granted all the quantities described in (i)–(iii), as desired.
It is essential to notice that each entry of x 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 X k +1 , one writes the value of X k +1 − ρC , then the next update is simply
X k +2 = (cid:2)
X k
+
+ φ k +1 e (cid:62) + fϕ (cid:62) k +1 (cid:3) + .
Finally, all the remaining updates in DROT only involve vector and scalar operations, and can thus be ﬁnished off with a simple auxiliary kernel.
E XPERIMENTAL RESULTS
In this section, we perform experiments to validate our method and to demonstrate its efﬁciency both in terms of accuracy and speed.
We focus on comparisons with the Sinkhorn method, as implemented in the POT toolbox 3 , 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/vienmai/drot .
We consider six instances of SK and its log-domain variant, called SK1–SK6 and SK-Log-1–SKLog-6, corresponding to η = 10 − 4 , 10 − 3 , 5 × 10 − 3 , 10 − 2 , 5 × 10 − 2 , 10 − 1 , in that order.
Given m and n , we generate source and target samples as x s and x t , whose entries are drawn from a 2D Gaussian distribution with randomized means and covariances ( µ s , Σ s ) and ( µ t , Σ t ) .
Here, µ s ∈ Rhas normal distributed entries, and Σ s
= A s A (cid:62) s , where A s ∈ R 2 × 2 is a matrix with random entries in [0 , 1] ; µ t ∈ R 2 has entries formed from N (5 , σ t ) for σ t > 0 , and Σ t ∈ R 2 × 2 is generated similarly to Σ s .
Given x s and x t , the cost matrix C represents pair-wise squared distance between samples: C ij = (cid:107) x s [ i ] − x t
[ j ] (cid:107) 22 for i ∈ { 1 , . . .
, m } and j ∈ { 1 , . . .
, n } and is normalized to have (cid:107) C (cid:107) ∞
= 1 .
The marginals p and q are set to p = 1 m /m and q = 1 n /n .
DROT always starts at X
= pq (cid:62) .
[Figure]
[Figure] P e r f
o r m a n c e r a t i o [Figure] P e
[Figure]
f o [Figure] a n c e r a t i [Figure]
Performance proﬁle
For each method, we evaluate for each (cid:15) >
the percentage of experiments for which | f ( X K ) − f ( X (cid:63) )
| /f
( X (cid:63) ) ≤ (cid:15) , where f ( X (cid:63) ) is the optimal value and f ( X K ) 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 (cid:15) given on the x -axis.
Each subplot shows the result on 100 random problems with m = n = 512 and ρ = 2 / ( m + n ) .
We can see that DROT is consistently more accurate and robust.
It reinforces that substantial care is needed to select the right η for SK.
The method is extremely vulnerable in single-precision and even in double-precision, an SK instance that seems to work in one setting can run 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).
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% conﬁdence interval, as a function of the dimensions.
Here, m and n range from 100 toand 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 so once every 10 iterations.
The result conﬁrms the efﬁciency of our kernel, and for very large problems the per-iteration runtimes of the two methods are almost identical.
It is also clear that DROT and SK can execute an iteration much faster than the log-domain
SK implementation. Finally, Figs 3(b)-(c) show the median times required for the total error (the sum of function gap and constraint violation) to reach different (cid:15) values.
Since SK struggles to ﬁnd 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.
4 We can see that DROT can quickly and consistently ﬁnd good approximates, while SK-Log is signiﬁcantly slower for small η and may not attain the desired accuracy with higher η . [Figure]
The experiments conﬁrm the strength of DROT both in speed and robustness: (i) each DROT iteration can be executed efﬁciently as in SK, and (ii) it can ﬁnd high accuracy solutions as a well-tuned log-domain SK {+does, but at a much faster speed. (cid:107) m = ϕ 
+} =
the [ · ] + is of the (1 / ( mn ) − 0 / ( m + n )) (cid:28) 0 .
This makes the ﬁrst few iterations of DROT identically zeros. To keep the canonical X 0 , we just simply set ρ 0 to 0 . 5 / log( m ) to reduce the warm-up period as m increases, which is certainly suboptimal for DROT performance.
This work was supported in part by the Knut and Alice Wallenberg Foundation, the Swedish Research Council and the Swedish Foundation for Strategic Research, and the Wallenberg AI, Autonomous Systems and Software Program (WASP). The computations were enabled by resources provided by the Swedish National Infrastructure for Computing (SNIC), partially funded by the Swedish Research Council through grant agreement no. 2018-
A E XTRA EXPERIMENTS
This section complements the experimental results in the main paper with further experiments for
DROT, SK and SK-Log in various settings [Figure] Performance proﬁle for MNIST We follow the construction in Cuturi (2013) and generate the cost matrix C ∈ R 784 × 784 as 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 a vector of intensity of 28 × 28 and normalize it to form the probability vectors p and q , respectively. Similarly to Figure 2, Figure 4 shows the fraction of problems that are successfully solved up to an accuracy level ε given on the x-axis. Here, each subplot is obtained by comparing 300 random pairs of images.
[Figure] Performance proﬁle for SK-Log Figure 5 shows the fraction of problems that are 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 a well-tuned log-domain SK method
[Figure]
Single runs performance Figure 6 depicts the actual performance of different algorithms for random OT instances used to produce the performance proﬁle in Figure 2. In addition, Figure 7 highlights the linear convergence phenomenon predicted by our theory. Sparsity of transportation By design, DROT efﬁciently ﬁnds sparse transport plans.
To illustrate this, we apply DROT to a color transfer problem between two images (see Blondel et al.
By doing so, we obtain a highly sparse plan (approximately 99% zeros), and in turn, a high-quality [Figure] [Figure] artiﬁcial image, visualized in Figure 8.
In the experiment, we quantize each image with KMeans to reduce the number of distinct colors to 750 .
We subsequently use DROT to estimate an optimal color transfer between color distributions of the two images [Figure]
B P ROOFS OF CONVERGENCE RATES
In order to establish the convergence rates for DROT, we ﬁrst collect some useful results associated with the mapping that underpins the DR splitting:
[Equation]
DR corresponds to ﬁnding a ﬁxed point to T , i.e. a point y (cid:63) :
y (cid:63) =
T ( y (cid:63) ) , or equivalently ﬁnding a point in the nullspace of G , where G ( y ) =
y − T ( y ) .
The operator T , and thereby also G , are ﬁrmly non-expansive (see e.g. Lions & Mercier (1979)).
That is, for all y, y (cid:48) ∈ dom ( T ) :
[Equation]
[Equation]
Let G (cid:63) be the set of all ﬁxed points to T , i.e. G (cid:63) =
{ y ∈ dom ( T ) | G
( y )
= 0 } .
Then (11) implies the following bound:
Lemma B.1.
Let y (cid:63) ∈ G (cid:63) and
y k +1 =
T ( y k ) , it holds that
[Equation]
Proof.
We have
[Equation]
Corollary B
[Equation]
Lemma B [Equation] 1 implies that
[Equation]
Proof.
The ﬁrst inequality follows directly from the non-negativity of the left-hand side of
Lemma B.1.
The latter inequality follows from:
[Equation]
Applying Lemma B.1 to the right-hand side with y (cid:63)
= P G (cid:63) ( y k ) yields the desired result.
B. P ROOF OF T HEOREM
Since DR is equivalent to ADMM up to changes of variables and swapping of the iteration order, it is expected that DR can attain the ergodic rate O (1 /k ) of ADMM, both for the objective gap and the constraint violation (He & Yuan, 2012; Beck, 2017).
However, mapping the convergence proof of ADMM to a speciﬁc instance of DR is tedious and error-prone.
We thus give a direct proof here instead.
We ﬁrst recall the DROT method:
[Equation]
Since X k +1 is the projection of Y k − ρC onto R m × n + , it holds that
[Equation]
Also, since X is a closed afﬁne subspace, we have
[Equation]
Next, for X ∈ R m × n + , Z ∈ X and U ∈ R m × n we deﬁne the function:
[Equation]
Let U k
+
=
( Y k − X k +1 )
/ρ , it holds that
[Equation]
By (14), we have (cid:104)
Y
k − X [Equation] [Equation]
X[-k +1 , Z k +1 , U k +1 ) − V ( X, Z, U k +1 ) ≤ 
 ρ (cid:104) Y k − X k +-]
[Equation]
, we thus arrive at
[Equation]
where the second step follows from (15).
Now, for any Y ∈ R m × n , we also have
[Equation]
Therefore, by adding both sides of (16) and (17)
[Equation]
we obtain
[Equation]
Since Z k +1 − X k +
=
Y k
+
− Y k
[Equation]
it holds that
[Equation]
Plugging (19) into (18) yields
[Equation]
which by dropping the last term on the right-hand side gives the ﬁrst claim in the theorem.
Let Y (cid:63) be a ﬁxed-point of the mapping T and let X (cid:63) be a solution of (1) deﬁned from Y (cid:63) in the manner of Lemma 2.1.
Invoking (20) with Z = X
= X (cid:63) ∈
R m × n + ∩ X and summing both sides of (20) over the iterations 0 , . .
, k − 1 gives
[Equation]
(21) where ¯ X k = (cid:80) ki
=
X i /k and ¯ Z k = (cid:80) ki
=1 Z i /k .
Since
[Equation]
, it follows that
[Equation]
Since Y is arbitrary in the preceding inequality
[Equation]
we must have that
[Equation]
from which we deduce that
[Equation]
By Lemma B
[Equation] [Equation]
we have
[Equation]
which implies that
[Equation]
Finally, plugging (24) into (23) and (22) yields
[Equation]
which concludes the proof.
B. P ROOF OF T HEOREM
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, we 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 implicitly, rather than explicitly via the multipliers deﬁned in the ADMM framework.
As a consequence, their proof needs to be adjusted to apply to our proposed algorithm.
To facilitate the analysis of the given splitting algorithm, we let e = 1 n and f = 1 m and consider the following equivalent LP:
[Equation]
subject to Ze = p
Z (cid:62) f = q X ≥
Z = X.
The optimality conditions can be written on the form
[Equation]
(26a)
[Equation]
(26b)
[Equation]
(26c)
[Equation]
(26d)
[Equation]
(26e)
[Equation] 26f [Equation] This means that set of optimal solutions, S (cid:63) , is a polyhedron
[Equation]
where M eq and M in 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 of the constraint violation given by:
[Equation]
For the ( X, Z ) -variables generated by (9), (26a), (26b), and (26c) will always hold due to the projections carried out in the subproblems.
As a consequence, only (26d), (26e), and (26f) will contribute to the bound (27).
In particular:
[Equation]
The following Lemma allows us to express the right-hand side of (28) in terms of Y k
+1 − Y k .
Lemma B.2.
Let X k , Z k and Y k be a sequence generated by (9) .
It then holds that
[Equation]
Proof.
The ﬁrst equality follows directly from the Y -update in (9).
Let
[Equation]
the Y -update in (10) yields
[Equation]
In addition [Equation] the X -update
[Equation]
can be combined with the ﬁrst equality in (30) to obtain
[Equation]
or equivalently
[Equation]
By (30)
[Equation]
the left-hand-side in (32) can be written as
[Equation]
Substituting this in (32) gives
[Equation]
which gives the second relation in the lemma.
Moreover, (31) and (33) implies that X k +1 and µ [Equation] 1 − ρ − 1 ( Y k +1 − Y k ) − C cannot be nonzero simultaneously. As a consequence, (cid:10)
[Equation] , µ [Equation] or (cid:104) C, X k +1 (cid:105) =
(cid:10) X k +1 , µ k +1 e (cid:62) + fν (cid:62)
k
+1 − ρ − 1 ( Y k +
[Equation]
Y k ) (cid:11)
[Equation]
ρ
[Equation]
Y [Equation] [Equation] Y k (cid:107)
[Equation]
By (30), ρ − 1 Z
[Equation]
1 + µ
[Equation]
fν
[Equation] [Equation]
X
[Equation] [Equation]
= ρ − 1 (2 Y
[Equation]
+1 − Y k ) , hence
[Equation]
[Equation]
Note that
[Equation]
Substituting (35) into (34), we get the third equality of the lemma:
[Equation]
By combining Lemma B.2 and (28), we obtain the following result:
Lemma B.3.
Let X k , Z k and Y k be a sequence generated by (9) and µ k , ν k be deﬁned by (29) .
Then, there is a γ > 0 such that:
[Equation]
Proof.
Let Y (cid:63) ∈ G (cid:63) , then Lemma B.1 asserts that
[Equation]
which implies that (cid:107) Y k (cid:107) ≤ M for some positive constant M .
In fact, { Y k } belongs to a ball centered at Y (cid:63) with the radius (cid:107) Y 0 − Y (cid:63) (cid:107) .
From the bound of (28) and Lemma B.2, we have dist
[Equation]
By letting γ = θ S (cid:63) (1 + ρ −
M + 1))
we obtain the desired result.
Recall that for DROT, prox ρf ( · ) =
P R m × n + ( · − ρC ) and prox ρg ( · ) =
P X ( · ) .
The following Lemma bridges the optimality conditions stated in (26) with G (cid:63) .
Lemma B.4.
If Y (cid:63) =
X (cid:63) + ρ ( µ (cid:63) e (cid:62) + fν (cid:63) (cid:62) ) , where X (cid:63) , µ (cid:63) , and ν (cid:63) satisfy (26) , then G ( Y (cid:63) )
= 0 .
Proof.
We need to prove that
[Equation]
is equal to zero.
Let us start with simplifying P R m × n + ( Y (cid:63) [Equation] ρC ) .
By complementary slackness,
[Equation]
and by dual feasibility
[Equation]
Consequentially, we have:
[Equation]
, which deﬁnes a normal cone of R m × n + ,
[Equation]
[Equation] [Equation] X
[Equation]
By using the deﬁnition of Y (cid:63) , and that ρ > 0 , this implies that
[Equation]
which corresponds to optimality condition of the projection
[Equation]
We thus have P X
[Equation]
+
[Equation]
But by the deﬁnition of Y (cid:63) ,
[Equation]
, which means that X (cid:63) =
P X (2 X (cid:63) − Y
[Equation]
) (following from the optimality condition of the projection).
Substituting all these quantities into the formula for G ( Y (cid:63) ) , we obtain G ( Y (cid:63) )
= 0 .
Lemma B.5.
There exists a constant c ≥ 1 such that dist [Equation] Y
[Equation]
Proof.
Let Y (cid:63) =
X (cid:63) + ρµ (cid:63) e (cid:62) + ρfν (cid:63) (cid:62) , where ( X (cid:63) , X (cid:63) , µ (cid:63) , ν (cid:63) ) =
P S (cid:63) (( X k , Z k , µ k , ν k )) .
According to Lemma B.4, we have that Y (cid:63) ∈ G (cid:63) .
Then dist
[Equation]
e
[Equation]
Since Y k
= X k
+ ρµ k e (cid:62) + ρfν (cid:62) k , it holds that dist( Y k , G (cid:63) ) ≤ (cid:107) X [Equation]
X
[Equation]
+
[Equation]
+
[Equation]
− ν
[Equation]
) (cid:62) (cid:107) ≤ (cid:107) X k − X (cid:63) (cid:107) + ρ (cid:107) e (cid:107)(cid:107) µ k − µ (cid:63) (cid:107) + ρ (cid:107) f
(cid:107)(cid:107) ν k − ν (cid:63) (cid:107) ≤ (cid:112) 1 + ρ 2 (cid:107) e (cid:107) 2 + ρ 2 (cid:107) f (cid:107) 2 (cid:112) (cid:107) X k − X (cid:63) (cid:107)
+ (cid:107) µ k − µ (cid:63)
[Equation]
where the last inequality follows from Lemma B
[Equation]
The fact that
[Equation]
follows from (12a) in Corollary B
[Equation]
[Equation]
B. P ROOF OF LINEAR CONVERGENCE
By combining Lemma B.5 with (12b) in Corollary B
[Equation] [Equation]
we obtain:
[Equation]
or equivalently
[Equation]
Letting r = c/ √ 1 + c 2 < 1 and iterating the inequality k times gives
[Equation]
Let Y (cid:63) :=
P G (cid:63) ( Y k ) .
Since prox ρf ( Y (cid:63) )
∈ X (cid:63) and X k +
=
prox ρf
( Y k ) , we have dist( X k +1 , X (cid:63) ) ≤ (cid:13)(cid:13) prox
[Equation]
( Y k ) − prox [Equation] Y [Equation]
[Equation]
≤ [Equation] k − [Equation]
= dist
[Equation]
Y
[Equation] [Equation] ) ≤ dist
[Equation] Y
[Equation] [Equation] (cid:63) )
r k , where the second inequality follows from the non-expansiveness of prox ρf ( · ) This completes the proof of Theorem 2.
C DR AND ITS CONNECTION TO ADMM
The Douglas-Rachford updates given in (5) can equivalently be formulated as:
[Equation]
If we let w k +1 = 1 ρ ( y k − x k +1 ) , we get.
[Equation]
By re-indexing the x and the w variables, we get:
[Equation]
This is exactly ADMM applied to the composite convex problem on the form (4).
D F ENCHEL -R OCKAFELLAR DUALITY OF OT
Recall that
[Equation]
and g
[Equation]
X ) = I { Y : A
[Equation]
Y [Equation] X
[Equation]
X
[Equation]
X [Equation] Recall also that the conjugate of the indicator function I C of a closed and convex set C is the support function σ C of the same set (Beck, 2017, Chapter 4).
For a non-empty afﬁne set X
= { Y : A ( Y ) =
b } , its support function can be computed as:
[Equation]
where X 0 is a point in X (Beck, 2017, Example 2.30).
By letting X 0 = pq (cid:62) ∈ X , it follows that
[Equation]
∗ [Equation] Since A ∗ : R m + n → R m × n : ( y, x ) (cid:55)→ ye (cid:62) + fx (cid:62) , any matrix U ∈ ran A ∗ must have the form µe (cid:62) + fν (cid:62) for some µ ∈ R m and ν ∈ R n .
The function g ∗ ( U ) can thus be evaluated as:
[Equation]
We also have
[Equation]
if − U − C ≤ 0 , + ∞ otherwise .
The Fenchel–Rockafellar dual of problem 4 in Lemma 2.1 thus becomes
[Equation]
subject to
[Equation]
which is identical to (3) upon replacing µ by − µ and ν by − ν .