Bringing regularized optimal transport to lightspeed: a splitting method adapted for GPUs
Authors: Jacob Lindbäck
Zesen Wang
Mikael Johansson
Venue: NeurIPS 2023 poster (2023)
Keywords: optimal transport
domain adaptation
splitting methods
gpu computations
TL;DR: We propose a flexible splitting method that can solve regularized OT problems fast on GPU.

Abstract: We present an efficient algorithm for regularized optimal transport. In contrast to
previous methods, we use the Douglas-Rachford splitting technique to develop
an efficient solver that can handle a broad class of regularizers. The algorithm
has strong global convergence guarantees, low per-iteration cost, and can exploit
GPU parallelization, making it considerably faster than the state-of-the-art for
many problems. We illustrate its competitiveness in several applications, including
domain adaptation and learning of generative models.

Bringing regularized optimal transport to lightspeed: a splitting method adapted for GPUs
Jacob Lindbäck EECS, KTH Stockholm, Sweden j jlindbac@kth.ss se
Zesen Wang EECS, KTH Stockholm, Sweden zesen@kth. se
Mikael Johansson EECS, KTH Stockholm, Sweden mi mikell @kth kakkkkkkkhh... se
Abstract
We present an efficient algorithm for regularized optimal transport. In contrast to previous methods, we use the Douglas-Rachtn Rachffor splitting technique to develop efficient solver that can handle broad class of regularizers. The algorithm has strong global convergence guarantees, low per-iteration cost, and can exploit GPU parallelization, making it considerably faster than the state-of-the-art for many problems. We illustrate its competitiveness in several applications, including domain adaptation and learning of generative models.
1 Introduction
Optimal transport (OT) an increasingly important tool for many MI. problems. It has proven successful in wide range of upplications, including domain adaptation [8, 19], learning generative models [2, 17], smooth ranking and sorting schemes [12], and long-tailed recognition [29]. The versatility of OT stems from the fact that it provides flexible framework for comparing probability distributions that incorporates the geometric structure of the underlying sample space [30]. Early on, practitioners were directed to LP solvers with poor scalability to solve OT problems, but this changed dramatically with the introduction of entropic regularization and the Sinkhorn algorithm for OT. Sinkhorn's simple parallelizable operations resolved computational bottleneck of OT and enabled the solution of large-scale problems at "lightspeed" [11]. Despite the computational advantages of entropic regularization, many applications depend sparse structured transportation plans, and the bias that entropic regularization introduces can significantly impact the performance of the downstream task being solved. In such settings, other structure-pre protmooti regularizers are typically preferred (see e.g. [10, 25, 19I). However, to the best of our knowledge, no framework exists that handles general regularizers for OT unifying way with numerical efficiency that comparable to Sinkhorn. To this end, we study OT for broad class of regularizers and develop an algorithm with similar computational performance Sinkhorn. Further, we benchmark our method against the -fate-of-the-att showing that our algorithm achieves up to 100x speed-up up for several applications.
For discrete probability distributions, solving regularized OT problems amounts to finding solution to the optimization problem
minimize XER
subject to X1, = p, X 1m =q.
Here p and q non-negative vectors that sum 1, C. is non-negative cost matrix, and the regularizer is a function that promotes structured solutions. To solve (1) fast, it is important that the large number of decision variables and the non-negativity constraint are handled in a memory-efficiee way. It is also crucial to manage the non-smoothness induced by the non-negativity constraints without altering the complexity of the algorithm. It is common
the transportation plan
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

practice to solve the OT problem by considering its dual since many regularization terms give rise to dual problems with structure that can be exploited [5]. Most notably, entropic regularization, i.c.
(xX) = E., Xij log Xij. where >> 0, enables deriving a simple alternating dual ascent scheme, which coincides with the well-known Kankhom-KKopp algorithm for doubly stochastic matrices [39]. An advantage of this scheme is that it is easy to parallelize and has a low per-iteration cost [1 l]. However, decreasing the regularization parameter c will slow down the convergence of the algorithm and ultimately render the algorithm numerically unstable, which particularly noticeable when low-precision arithmetic is used [27]. Conversely, increasing makes the transportation plan blurrier which can be problematic in applications when sparse solutions are desired [5].
For more general regularizers, it is often difficult to develop fast algorithms via the dual problem. However, if the regularizer is strongly convex, the dual (or semi-dual) will be smooth, and standard gradient-basee methods can be used. This is computationally tractable when the gradients of the dual can be computed efficiently [5]. Besides potentially expensive gradient computations, just as for Sinkhorn, lower regularization parameters will slow down the algorithm. Furthermore, regularizers that are not strongly convex cannot be handled in this framework. Therefore, we propose a different technique to solve (1) for a broad class of regularizers, including many non-smooth and non-strongly convex functions. Using the Douglas-Rachee Rachford splitting technique, we derive an algorithm with strong theoretical guarantees, that solves range of practical problems rapidly and accurately. In particular, it efficiently handles regularizers that promote sparse transportation plans, in contrast to entropic regularization.
1.1 Contributions
We make the following contributions:
We adapt the Douglas-Rachfore splitting technique to regularized optimal transport, extend- ing the recent DROT algorithm proposed in [27].
Focusing on a broad class of regularizers which includes quadratic regularization and group lasso as special cases, we demonstrate global convergence guarantees and an accelerated local linear convergence rate. This extends the available theoretical results for related OT solvers and captures the behavior of the iterates observed in practice.
We develop an efficient GPU implementation that produces high-quality optimal transport plans faster than the conventional Sinkhorn algorithm. We then show how the proposed solver can be used for domain adaption and produce solutions of better quality up 100 times faster than the seate-of-the-aee when implemented on GPU.
1.2
Related Work
The Sinkhorn algorithm is arguably one of the most popular OT solvers. and it can often find approximate solutions fast, even for large-scale OT problems. Many extensions and improvements have been proposed, including variations with improved numerical stability [36] and memory-efficiee versions based on kernel operations [15]. However, to our best knowledge, no framework exists for general regularizers that enjoy comparable computational properties. As an example, standard approach for group-lasso regularization is to linearize the regularization term and use Sinkhorn iteratively [9]. Although this approach is fairly generic, it requires that the transportation plan recovered in every iteration, adding significant computational overhead. For strongly convex regularizers, such as quadratically regularized OT, several dual and semi-dual methods have been proposed, e.g. [5, 26], with stochastic extensions [37] and non-convex versions to deal with cardinality constraints [25]. However, these methods are significantly harder to parallelize, rendering them slower than Sinkhorn for larger problems [30]. Moreover, just as for Sinkhorn, both convergence rates and numerical stability properties deteriorate significantly with lower regularization parameters.
promising research direction that has recently attracted interest in the community is to consider splitting methods as an alternative to Sinkhorn for OT problems. For instance, an accelerated primal- dual method for OT and Barycenter problems was proposed in [7], proximal splitting for a particular OT discretization was explored in [28], and an algorithm for unregularized OT based on Douglas- Rachford splitting was developed in [27]. The convergence of splitting methods is well-studied even for general convex problems [3]. For
splitting, tight global linear convergence
2

---

rates can be derived under additional smoothness and strong convexity assumptions [18]. Global linear rates have also been established for certain classes of non non-smooth and non-strongly convex Douglas-Rachford
problems [41, 1J. Moreover, Douglas- Rachford splitting can benefit from the inherent sparsity of the problem at hand. For certain classes of problems, these algorithms can identify the correct sparsity pattern of the solution in finitely many iterations, after which stronger local convergence rate starts to dominate [23, 31J. This paper contributes to this line of research, by introducing a splitting method for regularized OT with strong theoretical guarantees and unprecedented computational performance.
Notation
For any X,Y Y E R we let (X,Y) == -r((TT))) and |||||= == XX, R and mxxx denote the set of m matrices with non-negative entries and non-positive entries respectively. We let (g denote the indicator function over a closed convex set S, ie is (x( = 0 if X E S and (x) xo if X & S, and the relative interior of set S is denoted relint S. The projection onto R=Xx is denoted [t]+ which sets all negative entries of the inputted matrix zero. The subdifferential of an extended-real valued function h is denoted @h( X), and its proximal operator is defined prox,p ph (x) = argmin, h(2) + 2p ll - æ|I? where p > 0
2 Regularized Optimal Transport
We consider regularized OT problems on the form (1) where the function h is sparsity promoting in the sense that its value does not increase if an element of X is set to zero.
Definition 2.1 (Sparsity promoting regularizers). : Rmxx R U {+oo] said to be sparsity promoting if. for any X E Rmx"" h(X) 2 h(X,) for every X, E IR with (X,)ij E {0, Xij b.
Notice that this function class does not necessarily induce sparsity. For instance h = 0, or h -1. || meet the conditions of Definition 2.1. Besides these two examples, the class of sparsity promoting functions include, but are not limited to, the following functions.
4(X) Epec [|l (group lasso OT)
4(X) Eis X X1jarcsinh(X,,, /B) (X2, + 82) (hypentropic regularization)
+(X) Eis Wiyh (irl where wij 2 0 (weighted E (regulaarizatin
h(X) Ensess Lx., (x)) (constrained OT)
Conic combinations of sparsity promoting functions are also sparsity promoting. Moreover, many regularized OT problems, such as Gini-regularized OT [32] and regularizers on the form A can be converted to only involve sparsity promoting regularizers. Note that the negative Shannon entropy Σ Xi log X.ş is not sparsity promoting.
In this paper, we will develop scalable algorithm for solving OT problems with regularizers. Besides strong theoretical guarantees, the algorithm numerically stable and can be implemented efficiently on a GPU. Our approach builds on the recently proposed DROT algorithm for unregularized OT [27]. We extend the DROT algorithm to regularized OT problems, improve the theoretical convergence guarantees, and develop an extended GPU kernel. The resulting algorithm solves regularized OT problems faster, with higher precision, and in a more generic fashion than the state-of-the-art. To explain the algorithm, we review the splitting method for OT developed in [27].
2.1 Douglas-Rachford splitting for OT
Rachford splitting is a technique for solving optimization problems on the form
minimize f(x) + g(x)
(2)
=ER"
using the fixed point update y T(m)), ), where
T (v) = y + prox (2prox,f of (2) v) prox,, (v)
(3)
3

---

and p > O is a stepsize parameter. If y" is fixed point of T then x² = prox, (y*) is a solution Douglas-Rachfore Rachford iterations are expressed as
(2). Often a third iterate, 2g, is introduced, and the Douglas-
#b+1 = pf (yk), 2k+11 = prox pg (2xx+1 yk)
Mk++ = + 2x++1 2k+1-
(4)
As long as and 9 in (2) are closed and convex (possibly extended-real valued), and an optimal solution exists, the DR-splitting method converges to a solution (see, e.g. Theorem 25.6 in [3]).
Under additional assumptions on f and g. the method has even stronger convergence guarantees [18].
For given problem that can be formulated on the form (2), there are typically many ways to partition the problem between ff and g. This must be done with care: a poor split can result an algorithm with a higher per-iteration cost than necessary or one which requires more iterations to converge than otherwise needed. often difficult achieve both simultaneously, which poses a challenging trade-off between the iterate complexity and the per-iteration cost of the resulting algorithm.
To facilitate the derivation of the algorithm, we let {X R=*x" X1,, p, X 1m q) and introduce the indicator fuunctions 'gm and (x, that correspond to the constraints of (1), Recognizing that the projection onto is tractable (by invoking recently derived formulas for matrix projections [4]) the authors of [27] proposed the splitting
f(X) (C, lapxa (X) and g(X) (x()
and demonstrated that the update (3) can be simplified to
Xx+1 = [Yx pC1+s Y++1 = X + OK+1 1,7 + 1m
(5)
where th and spk are vectors given via the iterative scheme:
Tk+1 X4+1 In
8k+1 X 1m
Mh+1 = T1+1/mm+m)
0k+1 = 02 kk+1
ax+1 = a2 - -k+11
bk+1 = bb Sk+1.
Φk+1 (aR: 2rk+1 + (2 Im) /n1
(k+1 (bg 28 + (2 1n)) /m.
(6)
The most expensive operations in (6) are the matrix- vector products in the r and s updates. Nonethe- less, these operations can efficiently be parallelized on a GPU. Further, notice that Y₁ can be eliminated by substituting it in the X-update of (5), giving
X = [ + 1 + 1m - pC1+.
Hence, the splitting scheme only needs to update the transportation plan X₃ and the ancillary vectors and scalars in (6).
3 DR-splitting for regularized OT
When extending the DR-splitting scheme the regularized setting, it is crucial incorporate h in one of the updates of (4) in a way that makes the resulting algorithm efficient. Specifically, if we let
f(x)= (C, X) + 'gmxn (X) h(X), and g(X) =x(x),
the first update of (4) reads:
X prox (Yk) prox Imax (Yo pC).
This update is efficient given that prox pl gu
(-) is easy to compute. This is indeed the case when
h is sparsity promoting. We present the result in Lemma 3.1.
Lemma 3.1. Let f(X) = XC.X) + 'Rmx" (X) + h(X) where h(X) is sparsity promoting, closed, convex and proper when X E IR Then prox (X) prox ph ([X pC]+).
Lemma 3.1 enables us to integrate the regularizer into the algorithm with only a minor update. Using the split of (7), the updated Douglas Rachffooord algorithm, which we will refer to as RDROT, follows:
X = proxph ([Yx pC]+), Ya+1 = Xx+1 + dx+11. + 1mp
(8)

---

Under the assumption that h is closed, convex, and proper over the feasible set of (I), Theorem 25.6 in [3] guarantees that RDROT converges solution of (1). Before we derive even stronger convergence properties of the algorithm, we describe a few specific instances in more detail.
Quadratic Regularization. Letting h(X)
| yields
X = (1 + pa) | pC
(9)
An attractive feature of quadratic regularization is that shares similar limiting properties as entropic regularization when used in OT-divergences [14]. In contrast to entropically regularized OT, letting - O does not lead to numerical instability in our method. Furthermore, the quadratic term makes the objective of (1) strongly convex which renders its optimal solution unique. This can be helpful when OT is used to define a loss function for e.g. training generative models [17] or adjusting for long-tailed label distributions [29], since it preserves differentiable (see Section 3.3 for details).
sparsity and renders the OT cost
Group Lasso Regularization. With h(X) = A E Egeg |1Xalle, where G is collection of disjoint index sets, the RDROT update becomes
X (Y p0]++ XK+10
X E G.
(10)
This regularizer has been used extensively for OT-based domain adaptation [10]. The rationale is that each unlabeled data point in the test (target) domain should only be coupled to data points of the training (source) domain that share the same label. This can be accomplished by organizing the data so that the rows of X correspond to data points in the training domain and the columns to points in the test domain, and then using sparsity-inducing group-lasso regularizer in (1). In this setting, ç a collection of sets, each specifying which columns correspond to particular label. Consequentially. this regularizer will promote solutions that map each data point in the test domain single label.
3.1 Rate convergence
It is well-known that DR-splitting finds an e-accurate solution in O(1/e) iterations for general convex problems [21]. This a significant improvement over Sinkhorn's iteration complexity of ((1/e²) [24]. but the Sinkhorn iterations are very cheap to execute. A key contribution of [27] was the development of a GPU kernel that executes the DR updates so quickly that the overall method is faster than Sinkhorn for many problem instances. They also demonstrated global linear convergence by exploiting the geometrical properties of LPs, but the convergence factors are very difficult to quantify for step-size tuning. As illustrated in Figure and discussed in more detail below, neither the global linear rate nor the 1 1/k-rate are globally tight. Instead, the stopping criterion decays as O(1/k) during the initial steps of the RDROT algorithm, but then begins to converge fast linear rate. Hence, existing convergence analyses that neglect this local acceleration are unsatisfactory.
Rositive entries (normalized)
10F
107
Suboptinality
E0ml
2000
350
250
400
1.30
2000
200
350
250
400
Iteration
Iteration
Figure 1: An illustration of the convergence behavior of the algorithm for quadratically regularized problem. The stopping criterion and the number of positive entries were computed in every iteration. Notice that the iteration number which the sparsity pattern stabilizes is aligned with when the linear convergence starts. In this example, the stepsize is chosen to make this effect more evident.
In this section, we will develop a convergence analysis for DR splitting on regularized OT problems that explains this observed behavior. Based previous work local convergence properties of
5

---

DR-splitting [23], we establish that our algorithm (8) identifies the true sparsity pattern of the solution in finitely many iterations. When the sparsity pattern identified, a fast linear rate typically starts to dominate over the sublinear rate. To derive these guarantees, we need one additional assumption on the optimal solution of (1).
Assumption 1. Let Y* be a fixed point of (3) such that Yn - Y and let X = pf (Y*) =
proxph ph ([Y* pC]+)) We then assume that:
1 (xr (
xx) relint of())).
(11)
Assumption can be seen as stricter version of the traditional optimality condition. It is a common assumption in analyses of active constraint identification that can be traced back to [6]. We consider this assumption to be fairly weak since it holds for most relevant OT problems, except for some very specific cost matrices. A more elaborate discussion is included in the supplementary material.
Under the regularity conditions of (11), we derive the following two convergence results.
Theorem 1. Let h(X) be sparsity promoting, convex, closed, and twice continuously differentiable for X > 0. If Assumption holds. then there a K 2 1 such that for all l 2 K, Xk, = 0 if and only if X3 O.
For our problem, this result implies that there is K after which the sparsity pattern does not change. To prove this, we first show that the sparsityprromotin and smoothness assumptions imply that f and 9 are partly smooth with respect two manifolds and then invoke results presented [23]. Moreover, Theorem I allows us to derive the following local convergence rate guarantees:
Theorem 2. Assume that the conditions stated in Theorem hold. If is locally polyhedral, then RDROT enjoys a local linear rate: When k 2 K, then
[Xx+1-X*M] S p [|YK - Y"lle T (0, 1).
Moreover, the optimal rate is independent of the stepsize p.
The regularizer h is locally polyhedral, for instance, in the unregularized setting, and when weighted l regularization is used (i.e. h(X)= -ii Wiyl Wiyl- where Wij 2 0). Therefore, one can expect sublinear rate until the correct sparsity pattern of the transportation plan is identified. From that point onward, a fast linear rate will dominate. The local rate holds for many other regularizers beyond locally polyhedral ones, which has also been observed in experiments (see, e.g., [23] and J. However, when the local polyhedrality assumption is relaxed, one must account for the Hessian information of the regularizer, which makes it difficult quantify a rate using this proof technique.
Stepsize selection The stepsize implicitly determines how the primal residual descent is balanced against the descent in the duality gap. When the cost has been normalized so that ||'l| = 1, stepsize that performs well regardless of h is p = 2(m + n) This stepsize was proposed for DR- splitting the unregularized problem in [27]. We use this stepsize in all our numerical experiments and observe that our method is consistently competitive without any additional tuning.
Initialization A reasonable initialization of the algorithm is to let X, = pg and du = 0m-- Pp O₂, However, when using the stepsize specified above, this will result in X 0mxn for I S k < N, where N = 1( n)). By initializing do = (3(m + ==))=1(1 + m//m + 1m)mm. and poo = (3(m + =1))-t(1 + n/(m + n)) one skips these first N iterations, resulting considerable speed-up the algorithm, especially for larger problems. The derivation is added to the supplementary material. There is potential for further improvements (see e.g. [40]), since - d yp are related to dual potentials. But for simplicity, use this strategy throughout the paper.
a
Stopping criteria One drawback of the Sinkhorn algorithm is that its theoretically justified subpoti- mality bounds are costly to compute from the generated iterates. This is usually addressed by only computing the primal residuals || X pll. and ||X;; 1, - all every M iterations and terminate when both residuals are less than a chosen tolerance. In contrast, our algorithm computes the primal residuals in every iteration (see r'K and 8; (6)), so evaluating the stopping criterion does not add any extra computations. Moreover, the iterates de and yo will, in the limit, be proportional to the optimal dual variables, and can hence be used to estimate the duality residual and duality gap every iteration. For more details, refer to the supplementary material.

---

3.2 GPU implementation
Similar to the approach developed in [27], the RDROT algorithm can be effectively parallelized on a GPU, as long as the regularizer is sparsity promoting and its prox-mapping is easy to evaluate. We have developed GPU kernels for RDROT with quadratic and group-lasso regularization, respectively'. Below, we provide an overview of how the DROT-kerne' is adapted for the regularized case. More details of the implementation are included in the supplementary material.
For the GPU implementation, the main computation time is spent on updating the X matrix as described in (8). Compared with DROT [27], the RDROT kernel is also required to evaluate the prox-mapping. Specifically, for the given regularizers, the following additional operations are needed:
Quadratic regularization: Each element of -x-pp]+ needs to be rescaled with a constant to update X (see (9)). This only requires minor change and results in practically the same per-iteration-cos as for the unregularized case.
Group-lasso regularization: The norm of each group, i.e Xa+1.g- 9 E G is required to the determine group-specitt scaling factor (see (10)). To this end, the following additional steps are required: (1) a reduction for gathering the square of the elements and computing the scale for the group: and (2) a broadcast to apply the scale X For small problem sizes, the reduction and the broadcast can be done within a thread block via its shared memory. For large problem sizes, when the elements in one group can not fit in a single thread block, the results of the reduction are stored in global memory, and an additional kernel function introduced apply the scale for the group. This overhead increases the per-iteration cost little, but we show in our experiments that it is still significantly faster compared to other methods.
The remaining parts of the algorithm are for updating vectors, which have much smaller time complexity than the X matrix updates. These are handled with an ancillary kernel.
3.3 Backpropagation via Pytorch and Tensorflow wrappers
In many applications, e.g. [17, 35, 29], useful to use the solution of (1) to construct loss function. Most notably, the optimal value of (1) is directly related to the Wasserstein distance between p and g. which is a natural measure of the proximity between distributions with many advantagcous properties (see [30] for extensive treatment). When training generative models, for example, the objective is to transform a simple distribution into one that resembles the data distribution. It is therefore natural to use the Wasserstein distance to construct loss functions, see e.g [2]. But in order to backpropagate through such loss functions, one must differentiate the optimal value of (1) with respect to the cost matrix C' (which is parameterized by both a data batch and generated batch of samples). Using Fenchel duality, we can show that the optimal solution X is a gradient (or Clarke subgradient) of the loss, regardless of the regularizer. By applying similar argument to the dual problem, one can show that p (#*: ; is a reasonable gradient estimate of the OT cost with respect RDROT for both the forward and backward pass and
to the marginals (): Hence, can do not need to lean memory eeemory-epeessvve unrolling schemes or computationally expensive implicit differentiation techniques. We include further details on the derivation the supplementary material. To facilitate using RDROT for DL, we wrapped our GPU OT solvers as PyTorch and TensorFlow extensions that feature fast automatic differentiation. The extensions can easily be integrated with networks that parameterize cost matrices or marginals. In a forward pass, our extension solves the regularized OT problem using the cost matrix and marginals given the current iteration and stores the approximate solution X and p=' (Ø) ip") for the backward pass. The aforementioned technique subsequently used in the backward pass to efficiently obtain an approximate gradient, which is used to rescale the gradients of the weights of the remaining network.
4 Numerical experiments
All experiments mentioned below are performed using single NVIDIA Tesla V100 GPU.
'The implementation is publicly available at GitHub: http:://gtthbb.co//aaazeees
Regularized-DRO0
7

---

GPU performance of RDROT with quadratic regularization
To illustrate the numerical advantages of the algorithm, we compared the GPU implementation of RDROT for quadratically regularized OT (which we will refer to as QDROT) with a dual ascent method (L-BFGS), proposed in [5], which we implemented in PyTorch with all tensors loaded on the GPU. Although there are several other solvers for quadratic OT, such as the semi-dual method proposed in [5] and the semi-smooth Newton method from [26], we have omitted them from our experiments since their associated gradient/search direction computations scale poorly with data size and are difficult to parallelize. In our benchmark, we simulated 50 data sets of size 500 x 500 and 5000 5000, respectively. We generated cost matrices with Cj₃ = _ Xj 1 where x x; are simulated 2D Gaussian variables with random parameters specific for each data set. We used the quadratic regularization term ((m + n2) [| X QDROT, we included a = 0 for reference. We ran the algorithms for each dataset and regularization parameter until they reached accuracy level of 0.0001. For commpleteness, Py P implementation of QDROT. The results are displayed in Figure 2. in which it is clear that method is substantially faster than the L-BFGS method. Furthermore, when the GPU kernel is used, an even more significant speedup is achieved compared to the L--BGGS method. Further details and additional experiments with other dataset sizes included in the supplementary material.
with ranging between 0.0005 to 0.2. For
also benchmark
10
101
COROT ntarch
COROT itorch)
LBPGS
LBRGS
Run time (ms)
Run time (ms)
12
185
1P Extee 2018- 600 1202" 223p
Sp 3xx 5P 10 ue ax
Regularization
Regularization
(a) m = 500, n 500
(b) m 5000, n = 5000
Figure 2: Comparison between the RDROT (QDROT) using the GPU kernel, a PyTorch version, and an I-BFGS method applied to the dual for different quadratic regularization parameters run on GPU. 50 datasets of two different sizes were simulated.
Group Lasso and Domain Adaptation
OT is the backbone of many domain adaption techniques since it provides a natural framework to estimate correspondences between a source domain and a target domain. These correspondences can subsequently be used to align the respective domains to improve the test accuracy. Specifically, for given source data set {(x,, 3 )
J==1 and target set {x/ {x where {y') )== are data labels, we
define the cost matrix elements Cū = d(x), x/) using a positive definite function d. By computing a transportation plan X associated with (1) and cost matrix C, one can adapt each data point in the source set into the target domain via x p. ΣØ= X,,x(. It has been shown that group-lasso regularization improve the adaptation [10] (cf. the discussion in $3). To handle the regularizer, it is customary to linearize it and iteratively solve updated OT problems with Sinkhorn [9. 10]. The current implementation of group-lasso OT in the Python OT package POT uses this approach [16]. Besides the computational cost of iteratively reconstructing transportation plans, entropic regularization tends shrink the support of the adapted training data and hence underestimate the true spread in the target domain. This is thoroughly discussed in [15] and illustrated in Figure 3 where we adapt data sets with different regularizers. Of course, the shrinkage can partially be addressed by decreasing the regularization parameter, but this tends to slow down that algorithm and may even lead to numerical instability. To shed some light on this trade-off, we simulated 50 datasets with two features, 1500 training samples, and 1000 test samples. The target domain was simulated via random affine transformation of the source domain. Each had 2 unique labels that were uniformly distributed among the instances. The labels in the test set were only used for validation. The cost
Ciy- x
x²] was used and normalized so that ||C||oo 1. To compare the performance of
8

---

Before adaptation
Miter adaptation (Entropic)
Alter adaptation (Quadratic)
After adaptation (GL Only)
Figure 3: toy example illustrating the effect strongly convex regularizations have on the support of the resulting adapted domains. Pink and purple points mark the labeled training set and the gray points the test domain. Notice how the Sinkhorn-based method, as well as the Quadratic Regularized method transfers the source domain samples towards the center of the target domain, while our approach transfers the samples in a way that matches the true variation of the test data
the algorithms, we computed the time taken to reach tolerance of 10-1 and evaluted the Wasserstein distance between the adapted samples of each label and the corresponding samples of the test set. This means that the better the alignment, the lower the aggregated distance. For the method using entropic regularization, we varied the regularization parameters 0.001 to 10, and the group lasso regularization set to 0.001. The results presented in Table 1. Indeed, increasing the entropic regularization parameter to speed up the conditional gradient methods will worsen the alignment between the source domain and the target domain, just illustrated in Figure 3. We also note that our method consistently outperforms the alternative methods, both in terms of adaptation quality and in terms of time reach convergence. Additional experiments with other problem sizes and regularization are consistent with Table 1, see the supplementary. In addition to the experiments
Table 1: Performance benchmark of RDROT for Group-Lasso regularization (GLDROT) against the conditional gradient based Sinkhorn algorithm. Median and 10th and 90th percentile statistics are included. Our method is competitive in terms of both speed and adaption quality.
Reg.
Runtime (s) +
Agg. W2 dist. +
Method
Ent
GL
Median
q10
990
Median
910
990
GLSK [10]
1e-3
1e-3
3.77
3,68
8.90
0.311
0.0657
6,48
GLSK [10]
1e-1
1e-3
3.73
3.71
3.77
8.24
3.47
31.6
GLSK [10]
1et2
1e-3
1,86
1.86
1.88
52.8
10.9
311
GLDROT (ours)
1e-3
0.0619
0.0475
0.0771
0.0576
0.0262
0.182
GLDROT (ours)
5e-3
0.0384
0.0306
0.0474
0.0879
0.0358
0.274
above, we compare the per-iteration costs for several OT solvers in the supplementary material. These strengthen our claim further that our algorithm achieves "lightspeed" even for large problem sizes.
Training of generative models
As an illustration of the efficiency of the GPU version of the algorithm, we use our PyTorch wrapper to implement the Minibatch Energy Distance [35], loss function for training generative models via OT. The network consists of a generator, and a network that learns cost function, and the OT-based Minibatch Energy distance quantify the similarity between data batches and generated artificial samples. We performed experiments with image generation on the MNIST and CIFARI 10 datasets, with some of the generated samples shown in Figure The details of the experiments are included in the supplementary material. The resulting generator gives sharp images, and the training time was better or comparable with altermative frameworks, despite the low amount of regularization used.
5 Conclusions
The ability to solve large regularized optimal transport problems quickly has the potential to improve the performance of wide range applications, including domain adaption, learning of generative models, and more. In this paper, we have shown that the Rachford splitting technique can handle broad class of regularizers an integrated and computationally efficient manner. Using a carefully chosen problem split, we derived an iterative algorithm that converges rapidly and reliably,

---

3
3
0
7
2
6
8
O
8
0
/
/
O
8
5
8
8
9
7
/
/
9
/
2
5
4
5
7
/ 7
1
4
/
S
4
6
8
O
Figure 4: Generated samples by GANs trained on MNIST and CIFARTO datasets with the Minibatch Energy Distance. The original sizes of samples are (28,28) and (32.32), respectively.
and whose operations are readily parallelizable on GPUs. The performance of our algorithm is supported by both theoretical arguments and extensive empirical evaluations. In particular, the experiments demonstrate that the resulting algorithms can be up to two orders of magnitude faster than the current state-of oatte-of-th-arr We believe that this line of research has the potential to make regularized OT numerically tractable for a range of tasks beyond the ones discussed in this paper.
Acknowledgement
This work was supported in part by the Knut and Alice Wallenberg Foundation, the Swedish Research Council, and the Wallenberg AI, Autonomous Systems and Software Program (WASP). The compu- tations 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-05973.
10

---

References
[1] David Applegate, Oliver Hinder, Haihao Lu, and Miles Lubin. Faster first-order primal-dual methods for linear programming using restarts and sharpness. Mathematical Programming, pages 1-52, 2022.
[2] Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein generative adversarial networks. In International conference machine learning, pages 214-223. PMLR, 2017.
[3] Heinz H Bauschke, Patrick Combettes, et al. Convex analysis and monotone operator theory in Hilbert spaces, volume 408. Springer, 2011.
[4] Heinz H Bauschke, Shambhavi Singh, and Xianfu Wang. Projecting onto rectangular matrices with prescribed row and column sums. Fixed Point Theory and Algorithms for Sciences and Engineering, 2021 (1)1(1),1-20. 2021.
[5] Mathieu Blondel, Vivien Seguy, and Antoine Rolet. Smooth and sparse optimal transport. In International conference on artificial intelligence and statistics, pages 880-889. PMLR, 2018.
[6] James V Burke and Jorge Moré. On the identification of active constraints. SIAM Journal on Numerical Analysis, 25(5):1 197--121, 1988.
[7] Antonin Chambolle and Juan Pablo Contreras. Accelerated bregman primal-dual methods applied to optimal transport and wasserstein barycenter problems. SIAM Journal on Mathematics of Data Science, 4(4):1369-11
2022.
[8] Nicolas Courty, Rémi Flamary, Amaury Habrard, and Alain Rakotomamoniy. Joint distribution optimal transportation for domain adaptation. Advances in Neural Information Processing Systems, 30, 2017.
[9] Nicolas Courty, Rémi Flamary, and Devis Tuia. Domain adaptation with regularized optimal transport. In Joint European Conference on Machine Learning and Knowledge Discovery in Dutabuses, pages 274-289. Springer, 2014.
[10] Nicolas Courty, Rémi Flamary, Devis Tuia, and Alain Rakotomamoniy. Optimal transport for domain adaptation. IEEE Transactions on Pattern Analysis and Machine Intelligence, 1093))1888818533 2016.
[11] Marco Cuturi.
Sinkhorn distances: Lightspeed computation of optimal transport. Advances in neural information processing systems, 26, 2013.
[12] Marco Cuturi, Olivier Teboul, and Jean Jean-Philippe Vert. Differentiable ranking and sorting using optimal transport. In H. Wallach, H. Larochelle. A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
[13] Yann N Dauphin, Angela Fan, Michael Auli, and David Grangier. Language modeling with gated convolutional networks. In International conference on machine learning, pages 933-941. PMLR, 2017.
[14] Simone Di Marino and Augusto Gerolin. Optimal transport losses and Sinkhom algorithm with general convex regularization. arXiv preprint arXiv:2007.00977 2020.
[15] Jean Feydy, Thibault Séjourné, François-Xavee Xavier Vialard, Shun-ichi Amari, Alain Trouvé, and Gabriel Peyré. Interpolating between optimal transport and mmd using Sinkhorn divergences. In The 22nd International Conference on Artificial Intelligence and Statistics, pages 2681 3681-2690. PMLR, 2019.
[16] Rémi Flamary, Nicolas Courty, Alexandre Gramfort, Mokhtar Z. Alaya, Aurélie Boisbunon, Stanislas Chambon, Laetitia Chapel, Adrien Corenflos, Kilian Fatras, Nemo Fournier, Léo Gautheron, Nathalie T.H. Gayraud, Hicham Janati, Alain Rakotomamonjy, levgen Redko, Antoine Rolet, Antony Schutz, Vivien Seguy, Danica J. Sutherland, Romain Tavenard, Alexander Tong, and Titouan Vayer. Pot: Python optimal transport. Journal of Machine Learning Research, 22(78): 1-8, 2021.

---

[17] Aude Genevay, Gabriel Peyré, and Marco Cuturi. Learning generative models with Sinkhorn divergences. In International Conference on Artificial Intelligence and Statistics, pages 1608- 1617. PMLR, 2018.
[18] Pontus Giselsson and Stephen Boyd. Linear convergence and metric selection for Douglas- Rachford splitting and ADMM. IEEE Transactions on Automatic Control, 62(():533- -44, 2016.
[19] Xiang Gu, Yucheng Yang, Wei Zeng, Jian Sun, and Zongben Xu. Keypoint-guided optimal transport with applications in heterogeneous domain adaptation. Advances in Neural Information Processing Systems, 35:1 4972-14985. 2022.
[20] Warren L Hare and Adrian S Lewis. Identifying active constraints via partial smoothness and prox-regularity. Journal of Convex Analysis, 11(2):251266,, 2004.
[21] Bingsheng He and Xiaoming Yuan. On the O(1/n) convergence rate of the Bouglas-Rachfond alternating direction method. SIAM Journal on Numericul Analysis, 50(2):000770, 2012.
[22] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014.
[23] 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-9133 2017.
[24] 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, pages 3982-3991. PMLR, 2019.
[25] Tianlin Liu, Joan Puigcerver, and Mathieu Blondel. Sparsity-constraaned optimal transport. International Conference on Learning Representations, 2022.
[26] Dirk A Lorenz, Paul Manns, and Christian Meyer. Quadratically regularized optimal transport. Applied Mathematics & Optimization, 83(3):1919-149, 2021.
[27] Vien V Mai, Jacob Lindbäck, and Mikael Johansson. A fast and accurate splitting method for optimal transport: Analysis and implementation. International Conference on Learning Representations, 2022.
[28] Nicolas Papadakis, Gabriel Peyré, and Edouard Oudet. Optimal transport with proximal splitting. SIAM Journal on Imaging Sciences, 7(1):212-233 2014.
[29] Hanyu Peng, Mingming Sun, and Ping Li. Optimal transport for long-tailed recognition with learnable cost matrix. 2022.
[30] Gabriel Peyré, Marco Cuturi, et al. Computational optimal transport: With applications to data science. Foundations and Trendy@ in Machine Learning, 11(5-6):355-077, 2019.
[31] Clarice Poon and Jingwei Liang. Trajectory of alternating direction method of multipliers and adaptive acceleration. Advances in Neural Information Processing Systems, 32, 2019.
[32] Lucas Roberts, Leo Razoumov, Lin Su, and Yuyang Wang. Gini-regularized optimal transport with an application to spatio-temporal forecasting. arXiv preprint arXiv: 2017.
[33] R Tyrrell Rockafellar and Roger J-B Wets. Variational analysis, volume 317. Springer Science & Business Media, 2009.
[34] Tim Salimans and Durk P Kingma. Weight normalization: A simple reparameterizztioon to accelerate training of deep neural networks. Advances in neural information processing systems, 29, 2016.
[35] Tim Salimans, Han Zhang, Alec Radford, and Dimitris Metaxas. Improving GANs using optimal transport. International Conference Learning Representations, 2018.
12

---

[36] Bernhard Schmitzer. Stabilized sparse scaling algorithms for entropy regularized transport problems. SIAM Journal on Scientific Computing. 4143):1443-A1481, 2019.
[37] Vivien Seguy, Bharath Bhushan Damodaran, Rémi Flamary, Nicolas Courty, Antoine Rolet, and Mathieu Blondel. Large-scale optimal transport and mapping estimation. International Conference on Learning Representations, 2018.
[38] Wenling Shang, Kihyuk Sohn, Diogo Almeida, and Honglak Lee. Understanding and improving convolutional neural networks via concatenated rectified linear units. In international conference on machine learning, pages 2217-2225. PMLR, 2016.
[39] Richard Sinkhorn and Paul Knopp. Concerning nonnegative matrices and doubly stochastic matrices. Pacific Journal of Mathematics, 21(2):343-348, 1967.
[40] James Thornton and Marco Cuturi. Rethinking initialization of the Sinkhorn algorithm. In International Conference on Artificial Intelligence and Statistics, pages 8682-8698. PMLR, 2023.
[41] Sinong Wang and Ness Shroff. A new alternating direction method for linear programming. Advances in Neural Information Processing Systems, 30, 2017.
[42] Yao-Liang Yu. On decomposing the proximal map. Advances in neural information processing systems, 26, 2013.
13

---

A Theoretical results
A.1 Proof of Lemma 3.1
To prove Lemma 3.1, the following result is helpful:
Lemma A.1. Given that h is sparsity promoting, closed, convex, and 0. then Xij 0 implies that (proxps ph (X)))s O. Moreover, if Xij > 0, then (prox,s ph (X)). 2
Proof.
Let Y prox ph (X). By the definition of the proximal operator, every Z E Ikmx satisfies
h(Y) + 2p
| - - - h(Z) + 2p 11 X|
or
MY) + 2p
|| - X1| - 1| - XI S 1Z)
Moreover, since h sparsity promoting h(Z) h(Y), and holds that
1
(11X-
- - Z - x So.
(12)
2p
Let Z be equal to Y in all entries but the (i,j)th, which is set to zero. If Xi = 0, then
1
X F I xI xI = ーー Y
Y
2p
2p
and invoking (12) gives that Yú
0. If Xij > 0, on the other hand, then
I - X1
1 x12
1
(Vos (16s
X.,)² x3) )
2p
2p
and (12) gives
(Y, X1y)² X
which implies that Yú 2 The proof is complete.
Lemma A.I and Theorem 1 in [42] allows prove the composition lemma quite easily. We end this section with the proof.
Proof. (Lemma 3.1) For p(X) ' (X), it holds that
prox pp (X) [X]+
and
0p(X) = NR xn (X) (0)
for 2 0, where Nomxn (X) = (Y Rmxn YúJ = 0f Xij > 0, Yu S 0if
By Lemma A.1, if X,j 0, then (prox ph (X));j = 0. and hence (8p(X));; = (8p(prox ph (x))) Further, for positive entries X,₃ > 0, (8p(X))s (0), which implies that (Op(prox, (X))) 2 (Op(X)):j. Altogether, this means that
@p(prox, ph (X)) 2 0p(X).
Theorem in [42] then gives
prox proxpp+th ph (X) = prox (prox (X)) = proxph ([X]+)
and
prox of (X) = prox proxx(c.)+pptoo) (X)
prooxpp+oh (X pC)
proxph (IX - -
14

---

A.2 On the non-degeneracy condition of Asssumption
Assumption 1 is violated if Y is a fixed point to equation (3) and X proxps (Y*), but the regularity condition not fulfilled. That is, if
1 P (x-- xx- 0f(x)), while - =(xx (x+-xx))
xx) $ relint 0f(x)).
We can write these conditions more conveniently as
( r (x-xx) 0f(x*)) relint (x*))..
(13)
Let explore to what degree Assumption I can be seen restrictive when it is invoked in Theorem 1. To this end, we also assume that is twice continuously differentiable for > 0.
Recall that f (C, + aapx + h(X) and note that the inclusion (13) cannot be satisfied if
X > 0, since (f(x)) is a singleton. We therefore focus on the case when Of is set-valued, i.e. when X 0 for some (i,j). Since h(X) is differentiable for X O, only need to consider dt 0lgmen mm implies that (Ōtamxx rrelint Ōigmix )is 0 and
the zero entries. If X thus that mm m relint 0lgmen O. Consequentially, af(X*) \ relint 0f(x*) = C a
relint (x*) and
1 ((x y Y xx1 p
C. + g, where g E 0h(X*) relint Oh(X*).
By the convergence of the algorithm, equation (5) in the main document implies that there exist two vectors o² and yo* such that y - X* 1 + m This gives that the cost is on the form
c +1m
In other words, it is only cost matrices with specific (and rather particular) structure that violate Assumption For instance, if is smooth, e.g. Al| |12. Assumption only violated when the cost can be written as sum of two given rank-1 matrices. This is rarely met in real applications.
A.3 Proof of Theorem and 2
To prove the theorem need the notion of partial smoothness:
Definition A.1 (Partial Smoothness). lower semi-continuous function f IR. -> said to be partly smooth at x² E R" with respect to the C² (2-manifold M C IR" if:
(i) The restriction of f M: flM is C2 around x=,
(ii) (m*)*) = (par (xx))) where (par (f(x)))² denotes the orthogonal complement of the smallest subspace parallel to )x
(ii) (f(x) is continuous at x² relative to M.
Let X be solution the OT problem that fulfills Xx + X*, and let
M1 = (X E mARx Xij O, if X = 0, X, > 0 X > 0}
M3=
(14)
Recall that f(X) (C, 'gmx (X) h(X). and g gXX) 4x (X)).. Since h is assumed to be twice continuously differentiable over R by the conditions the theorem, partly smooth with respect to Mi Further, since is constant over M2. it is also partly smooth over M2. The optimality conditions of the regularized OT problem are
0 0E
Y* X
P
(15)
0 0g(x)) +
y
x
15

---

Moreover, since 0g = Nx, and X' is an affine subspace, 0g = relint dg. Therefore, Assumption
extends the following the non-degeneracy condition:
0 E relint 0f(x*)
p Y y"- X*
(16)
0 relint 0g(x") + Y X
With the partial smoothness and the non-degeneracy condition established, have all the building blocks needed for proving Theorem and 2.
Proof. (Theorem 1 and 2) By invoking the first order optimality condition associated with the first update in the DR- -plliting scheme in (4), we get:
Xk+1 proxp/ (Yk) = 0E 1 X - Yv 0f(x++1)
or
1( Yx -xx11) E @f((x+1))
(17)
Similarly, the second update in (4) associated with the inclusion:
Zx+1 -prox (2X Yn) = 0E 1 p
Zu+1 2X + .vs) 0g((x++)
or equivalently
(V--Xxt)) Xx+1 -1 P
1
1 1 Yx+ Yc+1 -1
(vet1-Va) E c((x+t)).
(18)
where we used the third equation of the DR-update, i.e Y841 - Ys Zk+1 X Notice
how (17) and (18) relate closely the optimality conditions (15). Now consider the function
F(X, 2) = f(X) + (()) p (X- - Y, X*). Since X E 1R and Z E X for all k 2 1, it holds that F(X++1, Zk+1) = (C, (C+1) + h(X++1) p (Xx+1 ZK+1, Y X*).
Since f + g are convex, closed, and attains minimizer, Theorem 25.6 in [3] can be used establish that Y + Y*, when k- + oc, where Y* is a fixed point of the DR-update. In particular,
|Y++1 Yell = |l(Y++1 - Y - (Y= Y S |Y - Y + I Y - 0,
which can be used to deduce that the third term of F tends zero + 0o:
(XR+1 ZE+I, Y X*) < 2k++ X ||||Y* - X = || - 0.
Consequentially,
F(XE+I, Zx-1) - x(x) +g(*))
(19)
As h is assumed be convex, proper, and lower semicontinuous, have that h is subdifferentially continuous its domain (including at X*) [33]. This gives
dist o, 0x F(Xe+1 Zo41) (044)))12444)1))
= dist 0, 0f(X++1) - ( -x=xx))}
S p - - -Y-) - -(X++1 - -X"1l
1(Yrs
$ | - -Y" - X - x"I
< 2p
" - y
0.
Here, the first inequality follows from that dist(r, S) inf [|| U|ا : S < | yll, for all y E S. The second is the Cauchy Schwarz inequality. and the third inequality follows from the non-expansiveness of proximal operators:
Xx+1 X
||roo (Yx) prox,s (Y") IYx Y
16

---

Similarly,
dist 0, 0z F2F+(X+1 Zot1)
dist 0, 0g(XX= + - Y -x-)
1
s [[xx+1 X*) -Yxx - (Yk+1 Yx)I
$ 2p '|16s y + '11| Ykll
+ 0.
Hence, dist 0, OF OF(X++1 k+1 2x+1)
0. This together with (16) and (19), and the fact that f is
partially smooth at X with respect to M₁. we apply Theorem 5.3 in [20] which proves the theorem. To prove Theorem 2, one can directly invoke Theorem 5.6 [23].
B Initialization
When using the initialization Xo pqd Φu Om and (Pu On, and the default stepsize, the updates in (6) gives that
X₁ 0mxn
Φs (k + 1)n-)(pp (p + (m + 1)
Pk (k: + 1)m=>lm- (q + (m + n)
when k = 1, .....N. N, where N = min Cijmn(m + (mp; + ngj + 1) 17 O(max(m. (m. n)).
When the cost is normalized, to get the simple initialization strategy: Xo 0mxn. dg = (3(m + (3(m+n)) +n) (1 m/ + n)) and SPo = (3(m + n)) (1 + n/ (m + 1,, that skips these first N iterations.
the rough approximation, Ct 1, pt m and qj n
C Derivation of Dual and Stopping criterion
A simple calculation reveals that g(X) (x(X) has the Fenchel conjugate
( (U. T) if = n1ll + 1
g (()
00,
otherwise.
n1l + Imv for which ("( p_ / + 9 Tv. Further, the conjugate of f(X) (C, + mpxx (X) + h(X) can be expressed as
Hence, any feasible dual variable is the form
(*) sup (U, X) - f(x)
x
sup (U C) - h(X)
X20
sup (UU C]+, M(X)
X20
S sup (IU Ct+, X) - -XX)
hr" h( - C]+)-
Moreover, since h is sparsity promoting,
h" ([U C1+ sup (UU C]+, X) - 1(X)
x
S sup (UU - (x)+ - -((x++)
sup (U C),,) - -N(X)
X20
f(),,,
17

---

meaning that (²( h² (U C]+). This gives the dual problem
maximize U=p1, +Imu
--((-0) - g² -
(1-1u1 1 (20)
To relate the iterates the optimality condition, define XR+1 proxps (Yk), Zx+1
proxpy pg (2Xx+1 Yk), and Us (Yk-1 -X4)/p and observe that
(Z++1 Xk+1) p (X Yx 0b 1m -K++1 (p (o=)) 1m (p-p))pp)
By theorem 25.6 in (3], 2x Xx+1 O, and Uk+1 U where U solution to (20). Hence (p (p'10x)1n 1x + mm (p Pp*) + U* meaning that p= Øx and o " will correspond to μ and v in
(p)
(20) respectively. Using that ræ and 5ę in (6) are directly related the primal residual, we can derive the following stopping criteria:
Tprimal == max(I|Xx1,, pll. ||x, 1m all) max([|r=||||||4/
gap == Xx) + h(Xx)) -p- pp' Φk + +xTxxx) (xx) + (0' + 1msk - pp1++)
-(C (p_ 6x + + (x) + (Xk, ou. Pp).
where AG denotes the duality gap deviation compared to the unregularized problem. Further, notice that by letting X = [ + 1x 1, + 1m - pCC]+, we can express [ + 1m - -pC]+ [X&+1 Xx which facilitates evaluating the duality gap. Notice that h" may not be finite. such settings, the dual is maximized over its effective domain, resulting in the addition of dual constraints. To account for these, one must include a dual residual that measures the constraint violation. We include some examples of stopping criteria in the following table.
Table 2: Termination criteria for selection of regularizers. Here we let || denote an arbitrary norm that is sparsity promoting.
h
A
Vduul
0
0 [I[(1 + x)xx+1 - Xx]+
[ Xul+1le
|
|
0
prox, ([8++1 - Xe]+) l|r
D Backpropagation through regularized OT-costs
We let h. h + 'nm= 1x. and define the OT-cost follows: OT^(C) infx (C,X) hx()).
Note that OT cost is equal to the optimal value of the OT problem (1), i.e. X a solution to (1), then (C) hes as (C) = - (-C),, and hence 7i*( (-C) + T((*)) = (-C, X*). Invoking subdifferential
+ X*) (X*). The OT cost is also closely related to the Fenchel conjugate of
properties of the Fenchel conjugate gives that X E ah (- C). Further
0c(-OT
dch:( -X*.
In particular, whenever the solution unique, such as when strongly convex, then CC)T, X Therefore, using the transportation plan estimated via RDROT is reasonable gradient estimate of the OT cost.
E
GPU Implementation
E.1 GPU Architecture Basics
Before introducing the detailed design of the GPU kernels, we recall some basie properties of modern GPU architectures. Since the kernel design is based on CUDA, which is compatible with NVIDIA GPUs, we focus on the architecture of NVIDIA GPUs. Nevertheless, many programming concepts, hardware architecture decisions, and performance bottlenecks are shared with other GPU designs.
18

---

E.1.1 Thread Organization
The fundamental unit of computation in NVIDIA GPUs is the Streaming Multiprocessor (SM). It is responsible for executing the parallel computations and managing the resources on the GPU. Each SM consists of multiple CUDA cores, shared memory, registers, and other components that enable efficient parallel processing. However, the concepts introduced below also apply to other modern GPU architecture that supports the SIMT execution model (like the AMD and Intel GPUs).
The parallel computations are organized in threads, grouped according to the following hierarchy:
1. The thread is the basic unit of code. All threads execute the same code, but they are endowed with an ID that can be used to parameterize memory access and control decisions.
2. A block is group of threads that execute together on an SM in the GPU. Threads within block can cooperate and communicate using shared memory.
3. The grid refers to the entire set blocks that will executed on the GPU. It represents the overall organization of parallel computation.
The smallest unit of thread execution on an NVIDIA GPU is the warp. It consists of 32 threads that scheduled and executed together on single SM. All threads within warp execute the same instruction at the same time (but on different data, parameterized by the thread and block IDs). The reduction within the warp is highly optimized and efficient.
E.1.2 Memory Hierarchy
NVIDIA GPUs, such as the NVIDIA Tesla V100, have more than eight types of memory including global memory, shared memory, registers, texture memory, constant memory, etc. We only focus the first three memory types, since they are the most relevant to the kernel design in the next section.
1. On NVIDIA GPUs, the global memory refers to the main memory available on the GPU. It is high-speed memory space used for storing data and instructions that are accessible by all the CUDA cores in the GPU. The global memory is the entry for exchanging data with the system RAM (random access memory). The global memory lies in the highest level of the memory hierarchy and has the largest capacity and the slowest access speed. NVIDIA Tesla V V100 has 32GB of global memory. An important feature of global memory is coalesced memory access, which will be discussed below and used in Section 1.2.1.
2.
The shared memory is a low-latency memory space, physically located on the GPU chip. The shared memory is used for data sharing within block, reducing memory access times significantly compared with using the global memory. The access speed of the shared memory is roughly 100x faster than the uncached global memory. Other characteristies that is allocated for and shared between threads in the same block, and that its capacity is limited compared with the global memory. For example, NVIDIA Tesla V100 with compute capability 7.0, has 96KB of shared memory per SM and 80 SMs in total.
3.
The register memory of NVIDIA GPUs is the fastest memory with the lowest latency. Registers are private to each thread and are not shared among other threads or blocks. They are used for temporary data storage and computation in a single thread. The access speed for registers is faster than the shared memory, but not by orders of magnitude. Register memory also has limited capacity. NVIDIA Tesla V100 has a limit of 64K 32-bit registers per SM.
E.1.3 Coalesced Memory Access
An important feature of global memory management is coalesced memory access. It is a memory access pattern that maximizes memory bandwidth and improves memory performance in parallel computing, particularly on GPUs. It refers to the way threads in warp access memory in contiguous and aligned manner, minimizing memory transactions and maximizing data transfer efficiency.
Depending on the architecture of the GPUs, the device can load or write the global memory via 32-byte, 64-byte, or 128-byte transaction that is aligned with their sizes, respectively. the threads in warp access a continuous and aligned memory block with the size of 128-byte (this could be 32 4-byte single-precision floats, for example), only costs one transaction. the continuous memory block is misaligned with the 128-byte pattern, it costs two transactions. If the memory accessed
19

---

is strided, which means no two elements are in the same 128-byte memory block, it will cost 32 transactions, even though it accesses the same amount of memory as in the coalesced case.
E.2 GPU Kernel Design
Our GPU kernels are the enhanced and optimized versions of the code for unregularized OT that accompanies the paper [27], available on Github². In the following subsections, we describe our improvements of the basic kernel, and detail the extensions that we have made to implement quadratic and group-lasse regularization. We will also present run-time measurements to demonstrate that our kernels execute even faster than the "light-speed" per-iteration times of Sinkhorn.
E.2.1 General Optimization
Similar to the kernel described in [27], the main kernel designed the way that each thread block contains 64 threads and is responsible for updating a sub-matrix with size (64, 64) in X. However, for small problem sizes where the required number of blocks may be smaller than the number of SMs, this design may cause low utilization of the GPU since not all of the SMs are used. To improve the utilization for small problem sizes, we adapt the number of columns assigned to each block to the problem size, reducing the per-block work load and increasing the SM utilization.
Careful management of the global memory access is critical to the performance of CUDA applications. In the kernel from [27], uncoalesced memory access happens when the number of rows is not a multiple of 32. For example, when the number of rows is 127 and the block tries to update the second column of X, the threads in the block will access 64 floats. However, the continuous block is not aligned with the 128-byte block structure, because the first element of the second column is in the same block as the last 31 elements in the first column. In this case, accessing 64 will require three 128-byte transactions while it could be two 128-byte transactions in an optimal solution. This can lead to strange situations where decreasing the problem size by one increases the per-iteration runtime significantly (6-20%). To alleviate the problem and to make full use of the coalesced memory access, we updated the kernel from [27] to handle shifted columns that match the 128-byte memory blocks. An illustration of the update is in Figure where the old working area of a block from [27] is marked with "Planned Block Work Area" and the updated working of a block is marked with "Actual Block Work Area". The comparison of the per-iteration runtime before and after the optimization in Table shows that the update solves the problem of uncoalesced memory access, and that the implementation performs more predictably on various problem sizes.
floats
Table 3: Comparison of the per-iteration runtime before and after the update for coalesced memory access. Note that the focus here on the relative change the per-iteration runtime after decreasing the problem size by one. The improvement in the per-iteration runtime under the same problem size is caused by other optimization techniques.
Problem Size
Runtime Before Update
Runtime After Update
4096 x 4096
0.3577
0.3453
4095 x 4095
0.3812 (+6.58%)
0.3417 (-1.04%)
8192 x 8192
1.281
1.152
8191 x 8191
1.548 (+20.80%)
1.153 (+0.10%)
10240 x 10240
1.916
1.808
10239 x 10239
2.189 (+14.24%)
1.781 (-1.52%)
E.2.2 Quadratically Regularized DROT
Similar to the implementation in [27], the matrices X and C' are stored in the global memory in column-major order. For the problem sizes that we target, the global memory is the only one that can fit the matrices. In addition, the performance penalty of storing them in a slower memory is limited since each element only needs be loaded once for every update. As the update of X is done column by column, the auxillary variables á and y can be shared and reused among the threads in a block. To speedup the memory access, we load the required elements of e and & into shared memory the beginning of the main kernel.
"The implementation is open-source and available at https:/ggihabbcom/vieenniiiin
20

---

Figure illustrates how the update of one sub-column is done in a thread block of the quadratically regularized DROT kernel. More specifically, an update is comprised the following steps
(1) Load the corresponding elements of Φ and yp from the shared memory.
(2) Load the corresponding elements of X and C' from the global memory. Compute [Y pCoj]+ and C2j- Xij. For the quadratically regularized case, we can directly update Xiş with Vip after multiplying it with the scale 1/(1 + pa)).
(3) Perform in-warp reduction to compute the column sum and the objective value in the block.
Use atomic operations to accumulate the results to corresponding row and column sums.
(5) This step is only used by the group-lasso kernel, introduced in Section E.2.3.
Similar to the implementation by [27], the iterations divided into even iterations and odd iterations.
1. For the even iterations, the kernel updates Xg+1 as
X = proxph ([Xs + @x+11, + 1m - -p[]+) pC
2. For the odd iterations, it updates Xx+1 as
X = proxph ([X++1 + @x+1l, + Imppk+1)+)
In this way, the kernel only needs access the matrix C once every second iteration.
Except for the general optimizations introduced in Section E.2.1, the only difference from the basic kernel is step (2) which additionally multiplies the result with scalar. This minor change introduces negligible changes the per-iteration runtime.
E.2.3 Group-lasso Regularized DROT
The kernel for group-lasso regularized DROT is slightly more involved, but also follows the five-step procedure illustrated in Figure More specifically, performs the following steps:
(1) Load the corresponding elements of Φ and y from the shared memory.
(2) Load the corresponding elements of X and C from the global memory. Compute [Y,, pC and Ciy- Xay. Compared with the quadratically regularized DROT, the kernel cannot update X,s until the scale of the group is ready in step (5).
(3) Perform in- ta-p reduction. For the group-lasso regularized case, we also reduce the sum of [Y, - ,.,, to compute the norm of the group.
(4) Accumulate the results to the corresponding sum of the row, the sum of the column, and the of the squared elements within the group using atomic operations.
(5) Compute the scale for the group and apply to the computed [Y,, - -pCai++ (For large problem sizes, where a thread block can only access part of the group, the process of reduction and broadcast done in an additional kernel function).
+ so as to update Xij-
Compared with the quadratically regularized case, the group-lasso regularized DROT requires an additional round for reduction and broadcast within the group. Moreover, for large problem sizes, when the intermediate results for the group do fit in the shared memory, it inevitable to access the matrix C' in every iteration. Both factors lead to increase in the per-ite iierrtteaiion runtime of the Group-lasso regularized DROT, which ends up being roughly twice as long as that of the quadratically regularized DROT. However, as we will show in the next section, its per-iteration runtime is still comparable with the Sinkhorn- Mirkhorn-KKopp algorithm.
E.2.4 Per-iteration Runtime
The comparison of the per-iteration runtime among all methods is given in Figure 6. For QDROT and the Sinkhorn-Knoph algorithm, random samples are generated with problem sizes from 100 to 10000. For GLDROT, 12 random samples are generated with problem sizes from 100 to 10000 and with and classes, respectively. Note that the Sinkhorn-baan based implementation of group-lasso OT in POT [16] does not return the number of iterations, reports the per-iteration run-time. This makes
21

---

Sum of Group Norms Gja
Sums of Cols X?
(5)
Sums of Rows X,
Cost
Transport Plan
Planned Block Work Area
Actual Block Work Area
Loaded in Shared Memory
Update X with 16 Threads
Within-War Reduction
Figure GPU Kernel Design
it impossible estimate the corresponding run-time numbers for the Sinkhorn-based group-lasso
OT [11]. We have therefore excluded this algorithm here, but expect it to have larger (and possibly significantly larger) per-iteration runtime than the Sinkhorn-Knopp algorithm [11].
To make a fair comparison between the Sinkhorn- Knopp algorithm ours, we cover four variants of the Sinkhomn-Knon algorithm in the comparison. The implementation of Sinkhorn-Knop in POT [16] computes the primal residual once every 10 iterations, since this operation is relatively costly compared with the iteration itself. Moreover, to increase the numerical stability of the algorithm and to improve the quality of the results with smaller weights of the entropic regularizer, the Sinkhorn- Knopp algorithm is improved by moving the computation to log scale. The decision compute the primal residual or not, and to perform the computations in linear or logarithmic scale leads to the four distinct versions of the Sinkhorn- Sinkhorn-Knopp algorithm that evaluate. Since methods work directly on X we let them compute the objective value and the primal residual in every iteration.
The results show that QDROT implementation has roughly the same per-iteration runtime the standard Barkhorn-Knopp algorithm. The per-iteration runtime the Sinkhorn-Knopp algorithm in log scale is roughly 4x longer than QDROT and is roughly 2.5-4.5x longer than GLDROT.
F Additional Experiments
This section describes additional experiments with RDROT using quadratic and group-lasso regu- larization datasets of various sizes and using different hyperparameters. The results
22

---

Sirkhom
Runtime be iteration (ms)
Sikchom wio stopcing oit.
Sinkhom Oog)
Smihhmm (lgg we stopping cnt.
GLDROT cesses)
cassss)
2000
4000
5000
8000
10006
Problem size [m n)
Figure 6: Comparison of the per-iteration cost for OT solvers. We compare the Sinkhorn-Knopp Algorithm [11], and log domain variant (with and without stopping criterion computations every 10th iteration), against QDROT (ours): Quadratically regularized DROT, and GLDROT (ours): Group-lasso regularized DROT.
consistent with the experiments that could be included in the main document and demonstrate significant advantages of our method compared the state-of-the-art.
Quadratic regularization In the main manuscript, we compared solver with L-BFGS method applied to the dual problem. Two problem sizes were considered: m = 1000, n = 1000. and 2000, 3000. In Figure 2, rerun the benchmark on 4 additional problem sizes. Notice that results are consistent with those presented the paper. showing that our results generalize to more settings beyond the ones presented in the paper.
LEFGS
Run time (ms)
Bun time
Regulerizeiion
Regular cetion
(a) 100, 100
(b) 1000, 1000
OOAOT
OOROY
Run time Imd
Run time
1-122 14.00 1-112 2T 1018 24.00
3x30 Regular
(c) m 2000. 3000
(d) m 10000, n 10000
Figure 7: Additional experiments of the RDROT (QDROT) and the dual L-BFGS method comparison for different quadratic regularization parameters and problem sizes. 50 datasets of an additional problem sizes were simulated.
23

---

Group Lasso regularization To establish the computational advantages of RDROT applied to the Group-Lasso problem, we benchmarked it against Sinkhorn- based solver on simulated datasets of size = 1000, 1500, for selection of hyperparameters. It Table 4,5,6, we rerun the benchmark on more problem sizes and hyperparameters, to test if the results generalize. Indeed, GLDROTT consistently outperforms the state-of-the-an in all considered problems.
Table 4: GL experiments with
500 and m
500
Reg.
Runtime (8) +
Agg.
W dist. +
Method GLSK
Ent
GL.
Median
q10
g90
Median
q10
790
1e-3
1e-6
1.26
1.24
2.98
0.35
0,084
4.96
1e-3
5e-4
1.23
1.22
2.93
0.351
0.0841
4.96
1e-3
5e-2
1.35
1.23
3.57
0.377
0.092
4.96
1e-2
1e-6
1.21
1.09
2.63
1.22
0.668
5.48
1e-2
5e-4
1.18
1.07
2.61
1.22
0.668
5.48
1e-2
5e-2
1.2
1.08
2.63
1.22
0.668
5.57
1e-1
1e-6
1.04
1.04
1.06
7.83
3.27
30.6
1e-1
5e-4
1.05
1.05
1.07
7.83
3.27
30.6
1e-1
5e-2
1.05
1.04
1.06
7.83
3.27
30.6
1.
1e-6
1.03
0.517
1.04
45.3
9.74
212
1.
5e-4
1,02
0.512
1.03
45.3
9.74
212
1.
5e-2
1.03
0.514
1.03
45.3
9.74
212
GLDROT
1e-6
0.0924
0.0519
0.16
0.0618
0.028
0.309
5e-4
0.0553
0.0354
0.106
0.0801
0.0357
0.319
5e-2
0.0127
0.00889
0.0351
0.672
0.321
2.87
Table 5: m = 1000, n
1000
Reg.
Runtime (s)
Agg.
W2 dist.
Method
Ent
GL
Median
q10
q90
Median
910
990
GLSK
1e-3
1e-6
4.97
4.91
9.27
1.66
0.405
11
1e-3
5e-4
4.9
4.82
9.19
1.66
0.405
11
1e-3
5e-2
5,52
4.91
15.9
1.64
0.389
11
1e-2
1e-6
4.9
4.79
5.53
3.06
0.926
8.97
1e-2
5e-4
4.79
4.66
5.5
3.06
0.926
8.97
1e-2 5e-2
4.88
4.73
5.55
3.15
0.926
9,03
1e-1
1e-6
4.71
4.68
4.73
20.6
8.97
74.3
1e-1
5e-4
4.74
4.7
4.78
20.6
8.97
74.3
1e-1
5e-2
4.76
4.73
4.82
20.6
8.97
74.3
1.
1e-6
4.61
4.57
4.64
130
37.5
403
1. 5e-4
4.62
4.6
4.67
130
37.5
403
1. 5e-2
4.74
4.71
4.77
130
37.5
403
GLDROT
1e-6
0.109
0.0735
0.146
0.0999
0.0424
0.304
5e-4
0.0598
0.045
0.0763
0.138
0.0632
0.44
5e-2
0.0281
0.0232
0.0374
3.16
1.14
9.01

---

Table 6:
1000, n
1500
Reg.
Runtime (8)
Agg.
W2 dist.
Method
Ent
GL
Median
q10
g90
Median
g10
990
GLSK
1e-3
1e-6
3.82
3.75
9.03
0.311
0.0657
6.48
1e-3
5e-4
3.8
3.75
9
0.311
0.0657
6,48
1e-3
5e-2
3.8
3.77
10.6
0.345
0.0665
6.48
1e-2
1e-6
3.78
3.73
6.8
1.21
0.69
5,47
1e-2
5e-4
3.8
3.76
6.82
1.21
0.69
5.47
1e-2
5e-2
3,82
3.79
6.88
1.21
0.69
5,47
1e-1
1e-6
3.7
3.67
3.72
8.24
3.47
31.6
1e-1
5e-4
3,75
3.73
3.78
8.24
3.47
31.6
1e-1
5e-2
3.77
3.73
3.81
8.24
3.47
31.6
1. 1e-6
3.73
1.86
3.77
45.5
9.92
218
1. 5e-4
3.69
1.85
3.73
45.5
9.92
218
1.
5e-2
3.75
1.88
3.78
45.5
9.92
218
GLDROT
1e-6
0.113
0.0758
0.147
0.0475
0.0215
0.137
5e-4
0.0745
0.0549
0.0951
0.0529
0.0245
0.163
5e-2
0.0232
0.0178
0.0288
0.331
0.154
1.38
G Model specification of the generative Adversarial Model
We adopted similar network structures and used the same loss function as in [35] and performed several experiments with image generation based on MNIST and CIFARIO datasets.
To define the loss, we need the following sample-based OT-cost.
(X,, Y) = (Cxx(
+hM) { {(XXX, + h(M): M1, = m7 "mm M 1mm = -n- -n
= inf MERT. "
Here m and n are the batch sizes corresponding X and Y, and Cx,y is matrix with pairwise distances between the samples of X and Y. The mini-batch energy distance [35] is then given by
CA = We, Y) + We.. (X', Y) + We, Y') + W.,
-2, x') - 2Wa, (Y,)") Y').
Here, X. X' are two independent mini-batches from real data, while Y, Y' two independent mini-batches generated from the generator.
In the experiments, used the cosine similarity parameterize the cost, and replaced the OT solver used in [35] the Kankhorn-Knoph algorithm, by the PyTorch wrapper for our OT solver with the quadratic regularizer.
The model structures are adapted from the ones [35]. Weight normalization is used to construct the parameters of the models [34], and the activation functions are gated linear units [13] and concatenated ReLUs [38]. We train the model with the Adam optimizier [22] using an initial learning rate of x 10 with 31 0.5 and B2 0.999. The batch size is 1024 for the MNIST experiment and 2048 for the CIFARIO experiment. We update the generator three times for every discriminator update. As for the parameters of OT solver,
10 the stopping criterion and = 10-3 as the weight of the quadratic regularizer. The results for MNIST and CIFAR10 are detailed below.
Different combinations of hyperparameters, model structures and schedules for the weight of the regularizers and the stopping criterion (the OT solver can allowed to produce low-accuracy solutions at the beginning of the training, for example) may potentially improve the results of the generated samples, but since GANs not the focus of this paper, leave such refinements future work.
25

---

G.1 MNIST Architecture and Results
Operation
Activation
Kernel Size
Stride
Padding
Output Shape
Sample
Linear
GLU
[128 7
Reshape
[128,
Upsample x2
[188, 14, 14
2D Convolution
GLU
[5. 5]
same
[128. 14, 14
Upsample x2
[288, 28, 28
2D Convolution
GLU tanh
[5,
same
64, 28, 28
2D Convolution
5,5 5
same
1. 28, 28
Table 7: Generator Architecture for MNIST
Opertion
Activation
Kernel Size
Stride
Padding
Output Shape
2D Convolution
CReLU
5, 5
same
128, 28, 28
2D Convolution
CReLU
5, 5
1
same
128, 28, 28
2D Convolution
CReLU
5, 5
2
2
[256, 14, 14]
2D Convolution
CReLU
[5,5]
2
2
[256, 7,
Flatten
256 7
L. Normalization
256 7
Table 8: Discriminator Architecture for MNIST
3
5
O
7
6
8
O
Y
3
9
3
8
0
/
O
8
5
S
/
O
5 6
9
9
(
2
6
f 5 5 7 1 6
8
1
8
9
/
/
9
1
/
8
D
6
7 8
/ 4
2
9 1 5
/
?
5
4
/
7
7
?
4 8
6
/
1
4
4
6
7
O
7
7
Y 6
O
4 3
7
9
6
9
7
/ 8
7
&
6
4
8
/ 7 7 4
4
7
7
9
3
6
/
5
8
6
8
O
4
7
2
4
4
0
4
7
5
8
2
7
6 7
O
8
3-
3
6 > 7 5
6
7
3
/
3
6
0
Z
O
7
3
6
1
6
9
4
5 /
I
O
O
3
O
7 9 4
5
\
8
5
3
9
L
L
5
4
6
7
6
9
7 ) 6 9
/
4
0
4 4
3 O
3
5
2
O
Q
5
2 6
2
>
8
6 2 2
7
6
2
3
4
7
4
8
3
3
9
/
5
/
/
/ 2
6
4
4
2
7
5
6
2
6
5
4
3
7
2
3
5 7 4 4 9
8
O
/
6
/
2
6
/
/
4
8
9
7
6
/
S
6
3 7
8
O 0
0
O
3
4 7
1
3
0
3
3
9
6
2
6
8
7
6
6
4
4
9
2 4
/
4
3
4
0
3
0
4
&
0 /
6
0
9
3
Figure 8: Generated samples at 100th epoch for MNIST (Original size for each sample is 28x28 px)
The generated samples in Figure are clear and shape, and no mode collapse can be observed (digits from 0 to 9 with various variants can be found in the figure). Compared with the results in [17] (Figure 3), our results show higher quality.
26

---

G.2 CIFARIO Architecture and Results
Opertion
Activation
Kernel Size
Stride
Padding
Output Shape
Sample
100
Linear
GLU
[1024 4 4]
Reshape
[1024, 4
Upsample x2
[1024, 8, 8
2D Convolution
GLU
[5 5]
same
[512, 8,
Upsample x2
[512, 16, 16
2D Convolution
GLU
[5, 5]
same
256,, 16, 16)
Up Upsample x2
256, 32,
2D Convolution
GLU
[5, 5] 5
same
128. 32,32
2D Convolution
tanh
same
3. 32,
Table 9: Generator Architecture for CIFARI0
Opertion
Activation
Kernel Size
Stride
Padding
Output Shape
2D Convolution
CReLU
5, 5
1
same
256, 32, 32
2D Convolution
CReLU
5, 5
2
2
512, 16, 16
2D
Convolution
CReLU
5, 5
2
2
[1024, 8,
2D Convolution
CReLU
[5,5]
2
2
[2048, 4,4
Flatten
2048 4
L2 Normalization
2048 4-
Table 10: Discriminator Architecture for CIFAR10
Figure 9: Generated samples at 1900th epoch for CIFARIO (Orig- inal size for each sample is 32x32 px)
In Figure 9, most samples are recognizable, and the samples cover the 10 categories of the CIFAR10 dataset with enough diversity. Our generated samples have similar quality as the ones reported in [35] (Figure 4) which uses the Kankhorn-Knopp algorithm the OT solver in the mini-batch energy loss.
27


Reviews:
[{"review_id": "gyzDWlvSq0", "reviewer": "Reviewer_WrJb", "summary": "In previous work [23] (DROT), the Douglas-Rachford splitting was applied to solving unregularized optimal transport (OT). The idea is to split the original variable into two variables $X$ and $Z$, where $X\\ge 0$ and $Z$ prescribed row and column sums. The current paper extends this idea to regularized OT. Compared to previous work [23], there are some novelties such as the local linear rate of convergence (Section 3.1, based on the results of Liang et al. [19]) and differentiation (Section 3.3).", "questions": "What happens if we use the entropic regularizer in RDROT?", "limitations": "See weaknesses.", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The paper is well written in general. Compared to DROT [23], this paper provides a local linear rate of convergence and a differentiation result. I particularly like the latter, it's a nice result and is useful in deep learning applications. \nThe obtained algorithm with its GPU implementation is very efficient and has the potential to replace existing algorithms in many applications (though I have to say that this is rather an encouragement, because the same thing could have been said about DROT [23], yet after more than two years since its publication, it hasn't been adopted in any research paper (3 citations at the time of writing, including one from the current submission), I wonder why this is the case.", "weaknesses": "The main limitation of this paper lies in its algorithmic contributions: there is virtually none. Unfortunately this is a major issue. The extension of DROT to the regularized case is rather straightforward. This is a very nice extension of DROT, but it alone is a rather weak contribution. \n\n\nSome presentation issues:\n\n- You should cite Sinkhorn and Knopp [31] when first mentioning the Sinkhorn algorithm in the introduction. Curuti's work [20] should be cited only for its efficient implementation.\n\n- You should cite Bauschke et al. 2021 (Projecting onto rectangular matrices with prescribed row and column sums) at line 138 when mentioning the projection onto X. It seems to me that this projection is the heart of both DROT and your algorithm. \n\n- Regarding the definition 2.1 of \"Sparsity promoting regularizers\" on page 3: according to this definition, it seems that the entropic regularizer is also a sparsity promoting one. However, the current discussion seems to indicate that the entropic regularizer is not suited for sparsity. Thus more clarifications and discussions are needed here (or maybe just choose another name for the definition).\n\n- I think the details on the stopping conditions and backpropagation are important and should be presented in the main content rather than in the appendix. \n\nMinor typos at line 93: <X,X> instead o <X,Y> in the definition of the norm. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thoughtful review and constructive feedback. We are grateful for your positive words about our work but, quite naturally, disagree with the statement that the algorithmic contributions are limited, or that this should be reason enough to reject the paper. Let us elaborate.\n\nWe believe that an important message in the work by Mai et al. 2022 was to highlight that for most existing algorithms, the unregularized OT problem is memory bound. In other words, even though there are many algorithms for solving the OT class of linear programs with an excellent iteration complexity, they are not amendable to efficient implementations beyond toy examples. In the same spirit, it is our opinion that the main contribution by Mai et al. was not to simply apply DR splitting to OT but rather to do it in a way that enables a memory-efficient GPU implementation.\n\nA limitation of the work by Mai et al. 2022 is that they focused on the unregularized OT problem. Although they reported crisp transport plans with significantly reduced blur compared to SK solutions, such high-accuracy solutions may not be needed in scenarios where the underlying data is noisy. This limits the practical usefulness of the DROT algorithm.\nThe main contribution of this work has been to identify a broad class of regularizers that can be dealt with efficiently on GPUs. Such problems have many applications but, quite surprisingly, few (if any) efficient algorithms. On the contrary, the algorithm that we propose is surprisingly simple, yet amazingly efficient. We believe that there are very few NeurIPS contributions that can, as we do, present an algorithm that runs 100x faster than the state-of-the-art on a well-established problem class. Indeed, one may argue that such an achievement alone should make the paper very interesting to a significant part of our community, and therefore qualify it for publication. To the best of our knowledge, the class of sparsity promoting regularizers is novel in the context of optimal transport, and its usefulness has not been pointed out in the literature before.\n\nAs you, and the other reviewers, point out, these are not the only contributions of the paper. We also do provide global and local theoretical guarantees that apply to the sparsity-promoting regularizers. Further, we show in our experiments that our framework is readily applicable to a range of problems, including domain adaptation and learning of generative models. To facilitate for the practitioner, we also wrapped RDROT in the autograd frameworks PyTorch and TensorFlow.\n\nWe believe there has been a minor misunderstanding regarding the question on entropic regularization. Entropic regularization: $H(X) = \\sum_{ij} h(X_{ij})$ with $h(x) = x \\log x$, is not sparsity-promoting, since it is not closed. By extending its domain so that$h(0) = 0$, it becomes closed, but still not sparsity promoting as $h(1/e) = -1/e < h(0)$. Therefore, our framework, in its current form, cannot handle this particular regularizer. However, if entropy-regularized OT is of interest, Sinkhorn-based approaches do an excellent job estimating optimal transportation plans. What we are trying to accomplish in our work is to go beyond this regularization scheme. One interesting direction to explore further is using hyperentropic regularization to interpolate between quadratic regularization and entropic regularization, which can be handled by our method. \n\nFurther, we think your input on the presentation issues are fair -  we will add the citations you recommended and fix the typo. We considered your suggestion to move the stopping criterion from the appendix, but we found that the current organization of the manuscript gives a better flow of ideas, and is more readable.\n\nWe hope that our answers will make you reconsider your scores.\n\nThanks!\nThe authors\n\n**References**\n\nVien V Mai, Jacob Lindbäck, and Mikael Johansson. A fast and accurate splitting method\nfor optimal transport: Analysis and implementation. International Conference on Learning\nRepresentations, 2022"}, {"review_id": "Nx23UguZFF", "reviewer": "Reviewer_e6yJ", "summary": "The paper presents an extension of the Douglas-Rachford algorithm in [23] to regularised optimal transport. It describes conditions under which sparsity-inducing regularisations result in sparse solutions and convergence rates of the resulting estimates.\nImplementation on GPU, as well as gradient routines, are considered in detail. Some numerical experiments illustrate the performance of the method.", "questions": "My main point of concern is the unfair empirical comparison. Because of this, I cannot trust the paper's title claim of \"lightspeed\", given that its performance may very well be implementation specific rather than methodological. I would be willing to substantially strengthen my rating if this was solved. See also the other (more minor and easy to fix) weaknesses.\n\n\n#### Minor points:\n- The notation for the prox operator $\\mathrm{prox}_{\\rho f}$ is undefined in the paper.\n- $e$ and $f$ in Equation (6) are undefined. Please replace with $\\mathbf{1}_{m/n}$ as necessary.\n- Typos: I did not spot any in the main text, but spotted two in the supplementary (assumpion, opertion), there may be other ones.\n", "limitations": "NA", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "To the best of my knowledge, the introduction of convex regularisation as part of a DR algorithm is novel. It seems to bring several benefits, both computational (speed + stability) and theoretical (sparsity and convergence rate guarantees). \n\nThe article is comprehensive in its treatment, considering parallelisation, gradient computation, and algorithmic designs such as step-size selection, etc. \n\nWhile I have not had time to go through the supplementary material in detail (I hope to be able to do so later on), the theoretical analysis looks sound and reasonable.", "weaknesses": "The experimental comparison with [4] is unfair: the implementation of the proposed method is fully tailored,  while the choice for [4] is a generic and not optimised PyTorch implementation. To improve this, the same effort should be spent on developing a custom CUDA kernel for [4] as for RDROT or providing a generic PyTorch implementation for both, rather than optimising the proposed method only. **IF** it is impossible that [4] can be improved, this must be explicitly discussed and explained.\n\nAdditionally, no comparison with Sinkhorn for entropy-regularised OT (and its accelerated versions, see for example http://angkor.univ-mlv.fr/~vialard/wwwsrc/AndersonMultiMarginal.html and https://github.com/ott-jax/ott/blob/main/src/ott/solvers/linear/acceleration.py)\nis given. Arguably, this is not sparsity-inducing, but (i) the title and introduction directly refer to [10], and (ii) the smoothness of the transport plan is sometimes a desirable property. Again, the comparison should be done with the same amount of implementation effort (either both custom-made kernels, or both high-level). \n\nWhile Section 3.3 discusses gradient computation, the supplementary is limited to gradient w.r.t. the cost function. In general, people tend to also be, if not more, interested in the gradient w.r.t. the marginals (weights and positions). Some discussion on this would be welcome. I also believe that (for at least the quadratic regularisation) the optimal transport plan is not differentiable w.r.t. the OT inputs. A quick discussion of this in the supplementary would be welcome.\n\nFinally, in terms of literature review, I am surprised that https://epubs.siam.org/doi/10.1137/130920058 (published in 2014) is not mentioned given that all the methods introduced in [23] (published in 2022) can, as far as I understand, be found in it. This does not impact the novelty of the current work (as I am not aware of a regularised version thereof, but proper attribution of the DR algorithm should go tohttps://epubs.siam.org/doi/10.1137/130920058 rather than [23].\n\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are grateful for the time you invested in reviewing our paper. Your feedback means a lot to us - thank you!\n\nWe have rerun our experiments with a PyTorch version of RDROT. Naturally, this results in a performance drop for our algorithm, but the experiments still show that our algorithm is faster than the state-of-the-art. Please see Figures 1 and Figure 2 in the authors' rebuttal. With that said, since the main computational bottleneck in RDROT is the update $X \\leftarrow X -\\rho C$, it is difficult to fully utilize the GPU parallelization and efficient memory management without tailoring a kernel (that is, by using working blocks and warp reductions). Sinkhorn-based methods are different since the memory-intensive operations are matrix-vector multiplications, which libraries such as PyTorch and TensorFlow handle efficiently. If you think commenting on this will improve the paper, we add a short comment in the final version of the paper.\n\n“Lightspeed computations” in this context refers to Cuturi’s seminal paper on Sinkhorn’s algorithm for OT. Therefore, when we assert that we bring regularized OT to lightspeed, we mean that we manage to develop an algorithm for regularized OT that has similar, or better performance than Sinkhorn (on the entropically regularized OT problem). As Figure 2 in the Supplementary material, and Figure 2 in the PDF attached to the Author rebuttal suggest, the per iteration cost of our method is similar to that of Sinkhorn. This together with the improved iteration complexity, in our view, suffice to say that our approach estimate transportation plans with lightspeed computations. If you believe that this would improve our paper, then we could also include explicit figures of the wall-clock execution times of the two methods. \n\nWith this understanding, techniques that accelerate the practical converge of SK execute “faster than light speed”. Our focus in this paper is really on regularized OT problems, and not on the pure OT problem that Mai et al considered, nor on multimarginal problems that is done in the post that you shared. For regularized problems, we are already able to attain significant speed-ups relative to the state of the art. It is an interesting idea to apply Anderson Acceleration also to the RDROT iterations, but it would be memory intense and probably difficult to scale even to medium-size OT problems. \n\nTo differentiate with respect to the marginals, one can use a similar argument discussed in the paper, but on the dual problem. We elaborate on this in the Author's Rebuttal. To our knowledge, most sparsity-promoting regularizers will result in transportation plans that are differentiable at least almost everywhere - but the resulting Jacobians will be zero. Developing ways of working around this is an active area of research, e.g. Sahoo et al. 2023, and it would be interesting to explore how this further can be applied to this framework.\n\nWe appreciate you sharing the typos you’ve found in our manuscript. We will revise according to your comments when polishing the final version.\n\nThanks a lot!\n\nThe authors\n\n**References**\n\nSahoo SS, Paulus A, Vlastelica M, Musil V, Kuleshov V, Martius G. Backpropagation through combinatorial algorithms: Identity with projection works, International Conference on Learning Representations 2023\n"}, {"review_id": "GAfZeZXAQm", "reviewer": "Reviewer_PMSB", "summary": "The paper develops an approximation algorithm for solving the optimal transport (OT) problem with a general class of regularizers (named \"sparsity promoting\", but more precisely characterized by not penalizing sparse solutions). The Douglas-Rachford splitting algorithm is used, extending the  previously introduced DROT algorithm to handle regularization. After introducing the OT problem and the Douglas-Rachford algorithm, the authors define a class of regularizers which include quadratic and group lasso regularization.\nWith additional assumptions, the convergence rate of the algorithm is established via first a convergence to the correct support, and then linear convergence to the optimal solution. The authors state that the proposed algorithm converges to epsilon accuracy in 1/epsilon iterations, instead of 1/(epsilon^2) of previously known methods. A fast GPU implementation of the RDROT algorithm is briefly described, along with some other practical considerations. Two numerical experiments are performed on a) domain adaptation, and b) generative modeling.", "questions": "If at all feasible, a comparison to more Sinkhorn-based algorithms (also in simpler settings, where convergence is assessed) could be useful to place the work within the literature.", "limitations": "A discussion on the tradeoffs between the proposed algorithm and minimizing sinkhorn divergences would certainly help to correctly evaluate the proposed contribution.", "rating": 7, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "The paper is well written overall, with a clear introduction to the problem, proposed approach and contributions. The practical considerations are also well written, and the GPU implementation -- as well as integration in popular programming frameworks -- makes the contribution particularly useful to the community.\nThe experiments are clear and present a convincing case for RDROT.", "weaknesses": "The theoretical aspects could be made clearer: it is not evident how theorems 1 and 2 lead to a 1/epsilon convergence rate (as there is no rate for support identification). The relevance of assumption 1 to the OT problem should be expanded upon: if and where it has been used in the literature, and when it doesn't apply. \nIn particular, a comparison to other convergence guarantees existing in the literature for Sinkhorn divergences (e.g. [1]) would be very useful in assessing the applicability of the proposed method.\n\n[1] Near-linear time approximation algorithms for optimal transport via Sinkhorn iteration, Jason Altschuler, Jonathan Weed, Philippe Rigollet, 2017", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your questions and your feedback on our paper! We believe that we have been able to address all your concerns, as detailed below.\n\nWe agree that there is room for improvement to make the theory section clearer. The 1/epsilon convergence rate of DR-splitting is nothing that we derive, but we simply refer to the work by He et al (2012) . This is a general result for DR splitting that is derived by leveraging the firm non-expansiveness of the proximal operators. The result only demands that the OT problem is convex and closed and that the problem has a solution. Hence, the 1/k convergence follows for all closed and convex regularizers that are proper over the convex polytope. Further, although the results in He et al (2012) are for ergodic sequences, some of the results can be generalized to non-ergodic sequences, see, e.g., He (2015). We will clarify this in the final version of our document.  \n\nIn addition to such existing results, we establish that our algorithm identifies the correct sparsity structure in a finite number of iterations and that stronger (typically linear) rates dominate once the correct sparsity structure is identified. This means that we not only have a global $1/k$ rate that is competitive to many alternative algorithms, but that we also, in many cases, have considerably stronger local rates. These stronger results are established under the additional Assumption 1. As discussed in Section A.2 in the Supplementary material, Assumption 1 holds for most OT problems besides some very specific edge cases. We have never encountered such edge cases in practice, and believe that they are rare. You are right that this assumption is typically not imposed in other OT papers. However, in the context of computational OT, to the best of our knowledge, no other work leverages sparsity to derive theoretical guarantees. Papers in other areas that do, such as Liang et al. (2015) , typically make this or related assumptions.  We will clarify these points in the final version of the document.\n\nThe reference you mention in the review (Altchuler et al. (2017)), but also Lin et al (2019), which we refer to in the paper, derive $1/\\epsilon^2$ rates for Sinkhorn and Greenkhorn. These rates are significantly worse than the ergodic rates that hold for RDROT. However, an interesting aspect of their analysis is that they quantify the total variation errors (i.e. $\\ell_1$-norm) of the marginals in terms of the problem size. We have not (yet) attempted to derive such results. Nevertheless, our numerical results strongly suggest that the numerical advantages of RDROT do not diminish as the problem size is increased (at least up to problems of the sizes that fit the memory of our GPU card).\n\nWe hope our rebuttal will make you more confident that the community would benefit from learning about our contributions at the NeurIPS Conference 2023.\n\nThank you!\n\nThe authors\n\n**References**\n\n\nHe B, Yuan X. On the O(1/n) convergence rate of the Douglas–Rachford alternating direction method. SIAM Journal on Numerical Analysis. 2012 \n\nHe B, Yuan X. On non-ergodic convergence rate of Douglas–Rachford alternating direction method of multipliers. Numerische Mathematik. 2015.\n\nLiang J, Fadili J, Peyré G, Luke R. Activity identification and local linear convergence of Douglas–Rachford/ADMM under partial smoothness. InScale Space and Variational Methods in Computer Vision: 5th International Conference. 2015.\n\nAltschuler J, Weed J, Rigollet P, Near-linear time approximation algorithms for optimal transport via Sinkhorn iteration. 31st Conference on Neural Information Processing Systems, 2017.\n\nTianyi Lin, Nhat Ho, and Michael Jordan. On efficient optimal transport: An analysis of greedy and accelerated mirror descent algorithms.  International Conference on Machine Learning, 2019."}, {"review_id": "qv1F29MzZz", "reviewer": "Reviewer_4HzB", "summary": "This paper adapts the Douglas-Rachford splitting to solve a wide range of sparsely-regularized optimal transport problems efficiently using GPU-parallelizable operations. The contributions are as follows:\n\n1) Adapt the Douglas-Rachford splitting to handle regularized OT problems with sparsity-inducing penalties\n2) Prove global convergence of the method and prove an accelerated local linear rate once the support is identified\n3) Provide GPU kernels for well-known regularizers (quadratic, group lasso) that achieve low-cost per iteration\n\n", "questions": "l.143: \"since $Y_k$ can be eliminated...\": could you elaborate on what you mean? After multiple readings, this remains unclear.", "limitations": "No negative societal impact for this work.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "The paper is well-written, the literature review is exhaustive and sets the context for this work: the relevant references are correctly cited. Besides, the subject tackled by the paper is critical for a wider adoption of OT to large-scale problems. The supplementary materials do a great job at making the experiments reproducible and explaining the low-level details of the kernel implementation.", "weaknesses": "The area of improvement for this paper remains the experiments section. For benchmarking purposes, I would use the Benchopt tool: https://github.com/benchopt/benchopt. Benchopt offers reproducible benchmarks and convergence curves over multiple runs. A figure from Benchopt to compare RDROT to L-BFGS would be more impactful than figure 2.\n\nMoreoever, the authors should favour real-world datasets instead of simulated cost matrices. The quadratic regularization and the Group Lasso sections should both display examples on real-world datasets to prove the robustness of the proposed method.\n\n=========\n\nBesides, I list below a few typos I spotted in the paper:\n\n1) l.49: Douglas-Rachford... splitting is missing.\n\n2) l.93: $\\lvert \\lvert X \\rvert \\rvert_F = \\sqrt{\\langle X, X \\rangle}$.\n\n3) l.108: I would precise that $\\lambda > 0$.\n\n4) In equation 6, $f$ and $e$ are extracted from [23], without explicitly defining them. I would explicitly use $\\mathbb{1}_n$ for $e$ and $\\mathbb{1}_m$ for $f$.\n\n5) l.165: \"used\" twice.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the time and effort you invested in thoroughly reviewing our paper. Your feedback is very valuable to us, as it helps us improve our paper. We are grateful that you highlighted some typos in the manuscript -  we will make sure to revise it according to your comments in upcoming versions.\n\nRegarding your question on why $Y_k$ can be eliminated, this comes from the fact that $Y_k$ can be substituted in the $X$-update of equation (5). This results in that $X_{k+1}$ only depends on $X_k$, $\\phi_k$, and $\\varphi_k$. Further, since the $\\phi$-updates are only dependent on $X$, we don’t have to keep track of $Y$ anymore after doing this substitution. This is pivotal for the efficiency of the resulting algorithm since it enables us to only keep track of one large matrix instead of two (or even three). To improve the readability of the final manuscript, we will add some additional explanations here.\n\nUnfortunately, up to this day, there are no OT benchmarks included in BenchOpt that we can use. Of course, this also means that there is an opportunity for us to contribute to the BenchOpt effort. Specifically, it would’ve been nice to have customized a QuadOT benchmark and GL-OT benchmark for this paper, but due to time constraints of the rebuttal, we haven’t been able to prioritize this. We will look into adding BenchOpt experiments to our main repository, to make our results more convincing and reproducible, and reach out to the BenchOpt authors to investigate the possibility to have these benchmark problems merged with their efforts.\n\nWe appreciate you sharing your input on improvements to the paper. We will update the manuscript according to the typos you highlighted in your review.\n\nMany thanks!\n\n\nThe authors\n\n"}]

Author rebuttal:
Dear AC and reviewers,

Thank you for the time you invested in the peer-reviewing process. The input you provided has both been insightful and inspirational for us. 

The reviewers have recognized many strengths and novelties of our contributions, including:

- the generality of the algorithm for regularized OT (it is readily applicable to many regularization schemes and constraints). We achieve this without compromising the numerical performance or the theoretical guarantees (e6yJ, PMSB, 4HzB).
- by using a tailored GPU kernel, our algorithm leads to ~100 x speedup compared to the state-of-the-art for several problems (WrJb, e6yJ, PMSB, 4HzB).
- we establish that a global competitive 1/k rate that holds for all considered regularizers, and a local linear rate under additional, weak assumptions (WrJb, e6yJ, PMSB).

For clarity, when we state that we bring regularized OT to lightspeed, what we mean is that RDROT has computational advantages comparable with Cuturi’s Sinkhorn method. We provide theoretical, numerical, and computational justifications for this claim. Our method enjoys rates that are better than that of the state-of-the-art, its per-iteration cost is low, and it benefits from GPU parallelization. Since our algorithm depends on operations that are difficult to parallelize in high-level frameworks such as PyTorch, to fully utilize the computational potential of the GPU, we’ve developed a GPU kernel for our framework. During this rebuttal, however, we found that an RDROT version implemented in PyTorch is still competitive for many problems. We illustrate this in Figure 1 in the attached PDF, in which we compare two versions of QDROT (RDROT applied to quadratically regularized OT), to an LBFGS method. Although the torch version is almost 10 times slower, it is still significantly faster than the LBFGS method, which is among the more popular algorithms for this application.

To enable practitioners to use our frameworks for deep learning, we developed a PyTorch and a TensorFlow wrapper that feature automatic differentiation through regularized OT costs. Some reviewers showed interest in this, including WrJb and e6yJ, which we are very pleased about. However, we stress that there are many possible extensions here. For instance, to differentiate with respect to the marginals, one can use the estimated dual variables $\phi/\rho$, and $\varphi/\rho$ as gradient approximations (see PDF for more details). To differentiate with respect to the transportation plan itself, one must use different techniques since the derivative is typically zero almost everywhere. We believe that the techniques proposed by Sahoo et al. (2023) could be worth exploring in this framework.

The feedback provided by the reviewers will contribute towards an even stronger final manuscript. The questions and criticisms raised by the reviewers have been taken into careful consideration - we are confident that we addressed all major concerns and the vast majority of the more minor suggestions. We believe our paper - which has been improved further by this reviewing process - would be of great value for the community, and hope that you agree that it is worthy of being presented at the conference in December.

Thank you!


The authors

**References**
Sahoo SS, Paulus A, Vlastelica M, Musil V, Kuleshov V, Martius G. Backpropagation through combinatorial algorithms: Identity with projection works, International Conference on Learning Representations 2023

Meta-review: This paper extends the Douglas-Rachford (DR) splitting method to sparsity-regularized optimal transport (OT). Though the idea of this extension is straightforward from the earlier DROT work, the authors demonstrated that their GPU implementation can be orders of faster than existing methods for solving the same problem, which shows the superiority of the proposed method.
Decision: Accept (poster)