Title: Bringing regularized optimal transport to lightspeed: a splitting method adapted for GPUs
Venue: NeurIPS 2023
Accepted: yes
Keywords: optimal transport

domain adaptation

splitting methods

gpu computations
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

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.

1 Introduction

Optimal transport (OT) is an increasingly important tool for many ML problems. It has proven successful in a wide range of applications, 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 a 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 a 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 on sparse or 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-promoting regularizers are typically preferred (see e.g. [10, 25, 19]). However, to the best of our knowledge, no framework exists that handles general regularizers for OT in a unifying way with a numerical efficiency that is comparable to Sinkhorn. To this end, we study OT for a broad class of regularizers and develop an algorithm with similar computational performance as Sinkhorn. Further, we benchmark our method against the state-of-the-art, showing that our algorithm achieves up to 100x speed-up for several applications.

For discrete probability distributions, solving regularized OT problems amounts to finding a solution to the optimization problem

Here p and q are non-negative vectors that sum to 1 , C is a non-negative cost matrix, and the regularizer h 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 on the transportation plan are handled in a memory-efficient 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

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.e. h ( X ) = ϵ ∑ ij X ij log X ij , where ϵ > 0 , enables deriving a simple alternating dual ascent scheme, which coincides with the well-known Sinkhorn-Knopp 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 [11]. However, decreasing the regularization parameter ϵ will slow down the convergence of the algorithm and ultimately render the algorithm numerically unstable, which is 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-based 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-Rachford splitting technique, we derive an algorithm with strong theoretical guarantees, that solves a 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-Rachford splitting technique to regularized optimal transport, extending 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 to produce solutions of better quality up to 100 times faster than the state-of-the-art 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-efficient 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, a 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 is 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.

A 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 primaldual 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 DouglasRachford splitting was developed in [27]. The convergence of splitting methods is well-studied even for general convex problems [3]. For Douglas-Rachford splitting, tight global linear convergence

rates can be derived under additional smoothness and strong convexity assumptions [18]. Global linear rates have also been established for certain classes of non-smooth and non-strongly convex problems [41, 1]. 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 a stronger local convergence rate starts to dominate [23, 31]. 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 ∈ R m × n , we let ⟨ X,Y ⟩ := tr ( X ⊤ Y ) and ∥ X ∥ F := √ ⟨ X,X ⟩ . R m × n + and R m × n -denote the set of m × n matrices with non-negative entries and non-positive entries respectively. We let ι S denote the indicator function over a closed convex set S , i.e ι S ( X ) = 0 if X ∈ S and ι S ( X ) = ∞ if X / ∈ S , and the relative interior of a set S is denoted relint S . The projection onto R m × n + is denoted [ · ] + , which sets all negative entries of the inputted matrix to zero. The subdifferential of an extended-real valued function h is denoted ∂h ( X ) , and its proximal operator is defined prox ρh ( x ) = argmin z h ( z ) + 1 2 ρ ∥ z -x ∥ 2 , where ρ > 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) . h : R m × n → R ∪ { + ∞} is said to be sparsity promoting if, for any X ∈ R m × n , h ( X ) ≥ h ( X s ) for every X s ∈ R m × n with ( X s ) ij ∈ { 0 , X ij } .

Notice that this function class does not necessarily induce sparsity. For instance h = 0 , or h = ∥ · ∥ 2 F 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.

h ( X ) = ∑ ij X ij arcsinh( X ij /β ) -( X 2 ij + β 2 ) 1 / 2 (hypentropic regularization)

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 ∥ X -A ∥ 2 F , can be converted to only involve sparsity promoting regularizers. Note that the negative Shannon entropy ∑ X ij log X ij is not sparsity promoting.

In this paper, we will develop a scalable algorithm for solving OT problems with sparsity-promoting regularizers. Besides strong theoretical guarantees, the algorithm is 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

Douglas-Rachford splitting is a technique for solving optimization problems on the form

using the fixed point update y k +1 = T ( y k ) , where

and ρ > 0 is a stepsize parameter. If y ⋆ is a fixed point of T then x ⋆ = prox ρf ( y ⋆ ) is a solution to (2). Often a third iterate, z k , is introduced, and the Douglas-Rachford iterations are expressed as

As long as f and g 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 a given problem that can be formulated on the form (2), there are typically many ways to partition the problem between f and g . This must be done with care: a poor split can result in an algorithm with a higher per-iteration cost than necessary or one which requires more iterations to converge than otherwise needed. It is often difficult to 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 = { X ∈ R m × n : X 1 n = p, X ⊤ 1 m = q } and introduce the indicator functions ι R m × n + and ι X , that correspond to the constraints of (1). Recognizing that the projection onto X is tractable (by invoking recently derived formulas for matrix projections [4]) the authors of [27] proposed the splitting

and demonstrated that the update (3) can be simplified to

where ϕ k and φ k are vectors given via the iterative scheme:

The most expensive operations in (6) are the matrix-vector products in the r and s updates. Nonetheless, these operations can efficiently be parallelized on a GPU. Further, notice that Y k can be eliminated by substituting it in the X -update of (5), giving

Hence, the splitting scheme only needs to update the transportation plan X k and the ancillary vectors and scalars in (6).

3 DR-splitting for regularized OT

When extending the DR-splitting scheme to the regularized setting, it is crucial to incorporate h in one of the updates of (4) in a way that makes the resulting algorithm efficient. Specifically, if we let

the first update of (4) reads:

This update is efficient given that prox ρh + ι R m × n + ( · ) is easy to compute. This is indeed the case when h is a sparsity promoting. We present the result in Lemma 3.1.

Lemma 3.1. Let f ( X ) = ⟨ C, X ⟩ + ι R m × n + ( X ) + h ( X ) where h ( X ) is sparsity promoting, closed, convex and proper when X ∈ R m × n + . Then prox ρf ( X ) = prox ρh ([ X -ρC ] + ) .

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-Rachford algorithm, which we will refer to as RDROT, follows:

Under the assumption that h is closed, convex, and proper over the feasible set of (1), Theorem 25.6 in [3] guarantees that RDROT converges to a 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 ) = α 2 ∥ X ∥ 2 F yields

An attractive feature of quadratic regularization is that it shares similar limiting properties as entropic regularization when used in OT-divergences [14]. In contrast to entropically regularized OT, letting α → 0 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 some sparsity and renders the OT cost differentiable (see Section 3.3 for details).

Group Lasso Regularization. With h ( X ) = λ ∑ g ∈G ∥ X g ∥ F , where G is a collection of disjoint index sets, the RDROT update becomes

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 a sparsity-inducing group-lasso regularizer in (1). In this setting, G is a collection of sets, each specifying which columns correspond to a particular label. Consequentially, this regularizer will promote solutions that map each data point in the test domain to a single label.

3.1 Rate of convergence

It is well-known that DR-splitting finds an ϵ -accurate solution in O (1 /ϵ ) iterations for general convex problems [21]. This is a significant improvement over Sinkhorn's iteration complexity of O (1 /ϵ 2 ) [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 a global linear convergence by exploiting the geometrical properties of LPs, but the convergence factors are very difficult to quantify or use for step-size tuning. As illustrated in Figure 1, and discussed in more detail below, neither the global linear rate nor the 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 at a fast linear rate. Hence, existing convergence analyses that neglect this local acceleration are unsatisfactory.

(normalized)

100

101

Kaewidogn

10-1

#Positive entries(

10-3

10-1

10-5

10-7

102

0

50

100

150

200

250

300

350

400

0

50

100

150

200

250

300

350

400

Iteration

Iteration

In this section, we will develop a convergence analysis for DR splitting on regularized OT problems that explains this observed behavior. Based on previous work on local convergence properties of

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 is 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 Y k → Y ⋆ , and let X ⋆ = prox ρf ( Y ⋆ ) = prox ρh ([ Y ⋆ -ρC ] + ) . We then assume that:

Assumption 1 can be seen as a 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 1 holds, then there is a K ≥ 1 such that for all k ≥ K , X k, ij = 0 if and only if X ⋆ ij = 0 .

For our problem, this result implies that there is a K after which the sparsity pattern does not change. To prove this, we first show that the sparsity-promoting and smoothness assumptions imply that f and g are partly smooth with respect to two manifolds and then invoke results presented in [23]. Moreover, Theorem 1 allows us to derive the following local convergence rate guarantees:

Theorem 2. Assume that the conditions stated in Theorem 1 hold. If h is locally polyhedral, then RDROT enjoys a local linear rate: When k ≥ K , then

Moreover, the optimal rate is independent of the stepsize ρ .

The regularizer h is locally polyhedral, for instance, in the unregularized setting, and when weighted ℓ 1 regularization is used (i.e. h ( X ) = ∑ ij w ij | X ij | , where w ij ≥ 0 ). Therefore, one can expect a 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 [31]). However, when the local polyhedrality assumption is relaxed, one must account for the Hessian information of the regularizer, which makes it difficult to 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 ∥ C ∥ ∞ = 1 , a stepsize that performs well regardless of h is ρ = 2( m + n ) -1 . This stepsize was proposed for DRsplitting on 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 0 = pq ⊤ , and ϕ 0 = 0 m , φ 0 = 0 n . However, when using the stepsize specified above, this will result in X k = 0 m × n for 1 ≤ k ≤ N , where N = O (max( m,n )) . By initializing ϕ 0 = (3( m + n )) -1 (1 + m/ ( m + n )) 1 m , and φ 0 = (3( m + n )) -1 (1 + n/ ( m + n )) 1 n , one skips these first N iterations, resulting in a 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 ρ -1 ϕ k , ρ -1 φ k are related to dual potentials. But for simplicity, we use this strategy throughout the paper.

Stopping criteria One drawback of the Sinkhorn algorithm is that its theoretically justified subpotimality bounds are costly to compute from the generated iterates. This is usually addressed by only computing the primal residuals ∥ X k 1 m -p ∥ , and ∥ X ⊤ k 1 n -q ∥ 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 s k in (6)), so evaluating the stopping criterion does not add any extra computations. Moreover, the iterates ϕ k and φ k will, in the limit, be proportional to the optimal dual variables, and can hence be used to estimate the duality residual and duality gap in every iteration. For more details, we 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 1 . Below, we provide an overview of how the DROT-kernel 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 [ Y k -ρC ] + needs to be rescaled with a constant to update X k (see (9)). This only requires a minor change and results in practically the same per-iteration-cost as for the unregularized case.

Group-lasso regularization: The norm of each group, i.e ¯ X k +1 ,g , g ∈ G is required to the determine a group-specific 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 to ¯ X k +1 ,g . 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 is introduced to apply the scale for the group. This overhead increases the per-iteration cost a 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], it is useful to use the solution of (1) to construct a loss function. Most notably, the optimal value of (1) is directly related to the Wasserstein distance between p and q , which is a natural measure of the proximity between distributions with many advantageous properties (see [30] for an 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 a 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 a similar argument to the dual problem, one can show that ρ -1 ( ϕ ⋆ ; φ ⋆ ) is a reasonable gradient estimate of the OT cost with respect to the marginals ( p ; q ) . Hence, one can use RDROT for both the forward and backward pass and do not need to lean on memory-expensive unrolling schemes or computationally expensive implicit differentiation techniques. We include further details on the derivation in 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 at the current iteration and stores the approximate solution X ⋆ k and ρ -1 ( ϕ ⋆ ; φ ⋆ ) for the backward pass. The aforementioned technique is 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 a single NVIDIA Tesla V100 GPU.

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 × 500 and 5000 × 5000 , respectively. We generated cost matrices with C ij = 1 2 ∥ x i -x j ∥ 2 2 , where x i , x j are simulated 2D Gaussian variables with random parameters specific for each data set. We used the quadratic regularization term (( m + n ) α/ 2) ∥ X ∥ 2 F with α ranging between 0 . 0005 to 0 . 2 . For QDROT, we included α = 0 for reference. We ran the algorithms for each dataset and regularization parameter until they reached an accuracy level of 0 . 0001 . For commpleteness, we also benchmark a PyTorch implementation of QDROT. The results are displayed in Figure 2, in which it is clear that our 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-BFGS method. Further details and additional experiments with other dataset sizes are included in the supplementary material.

(a)

m

= 500

,

n

= 500

(b)

m

= 5000

,

n

= 5000

103

QDROT

103

QDROT

QDROT (torch)

QDROT (torch)

LBFGS

LBFGS

Run time (ms)

Run time (ms)

102

102

101

101

100

100

0

-0x-0x-0x9-0x-0x-0x5-0x-0x

Regularization

0

-0x-0x-0x9-0x-0x-0x5g-0x-0x

Regularization

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 a given source data set { ( x i s , y i ) } i m =1 and a target set { x j t } n j =1 , where { y i } i m =1 are data labels, we define the cost matrix elements C ij = d ( x i s , x j t ) 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 i s,a = 1 p i ∑ m j =1 X ij x j t . It has been shown that group-lasso regularization can 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 to 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 a random affine transformation of the source domain. Each set had 2 unique labels that were uniformly distributed among the instances. The labels in the test set were only used for validation. The cost C ij = 1 2 ∥ x i s -x j t ∥ 2 was used and normalized so that ∥ C ∥ ∞ = 1 . To compare the performance of

Beforeadaptation

Afteradaptation(Entropic)

Afteradaptation(Quadratic)

Afteradaptation(GLOnly)

the algorithms, we computed the time taken to reach a tolerance of 10 -4 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 was set to 0 . 001 . The results are 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 as 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 to reach convergence. Additional experiments with other problem sizes and regularization are consistent with Table 1, see the supplementary. In addition to the experiments

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], a loss function for training generative models via OT. The network consists of a generator, and a network that learns a cost function, and the OT-based Minibatch Energy distance to quantify the similarity between data batches and generated artificial samples. We performed experiments with image generation on the MNIST and CIFAR10 datasets, with some of the generated samples shown in Figure 4. 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 alternative 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 a wide range of applications, including domain adaption, learning of generative models, and more. In this paper, we have shown that the Douglas-Rachford splitting technique can handle a broad class of regularizers in an integrated and computationally efficient manner. Using a carefully chosen problem split, we derived an iterative algorithm that converges rapidly and reliably,

8

8

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-the-art. 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.

References

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.

Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein generative adversarial networks. In International conference on machine learning , pages 214-223. PMLR, 2017.

Heinz H Bauschke, Patrick L Combettes, et al. Convex analysis and monotone operator theory in Hilbert spaces , volume 408. Springer, 2011.

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-20, 2021.

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.

James V Burke and Jorge J Moré. On the identification of active constraints. SIAM Journal on Numerical Analysis , 25(5):1197-1211, 1988.

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-1395, 2022.

Nicolas Courty, Rémi Flamary, Amaury Habrard, and Alain Rakotomamonjy. Joint distribution optimal transportation for domain adaptation. Advances in Neural Information Processing Systems , 30, 2017.

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 Databases , pages 274-289. Springer, 2014.

Nicolas Courty, Rémi Flamary, Devis Tuia, and Alain Rakotomamonjy. Optimal transport for domain adaptation. IEEE Transactions on Pattern Analysis and Machine Intelligence , 39(9):1853-1865, 2016.

Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. Advances in neural information processing systems , 26, 2013.

Marco Cuturi, Olivier Teboul, and 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.

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.

Simone Di Marino and Augusto Gerolin. Optimal transport losses and Sinkhorn algorithm with general convex regularization. arXiv preprint arXiv:2007.00976 , 2020.

Jean Feydy, Thibault Séjourné, François-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-2690. PMLR, 2019.

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, Ievgen 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.

Aude Genevay, Gabriel Peyré, and Marco Cuturi. Learning generative models with Sinkhorn divergences. In International Conference on Artificial Intelligence and Statistics , pages 16081617. PMLR, 2018.

Pontus Giselsson and Stephen Boyd. Linear convergence and metric selection for DouglasRachford splitting and ADMM. IEEE Transactions on Automatic Control , 62(2):532-544, 2016.

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:14972-14985, 2022.

Warren L Hare and Adrian S Lewis. Identifying active constraints via partial smoothness and prox-regularity. Journal of Convex Analysis , 11(2):251-266, 2004.

Bingsheng He and Xiaoming Yuan. On the O(1/n) convergence rate of the Douglas-Rachford alternating direction method. SIAM Journal on Numerical Analysis , 50(2):700-709, 2012.

Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980 , 2014.

Jingwei Liang, Jalal Fadili, and Gabriel Peyré. Local convergence properties of DouglasRachford and alternating direction method of multipliers. Journal of Optimization Theory and Applications , 172(3):874-913, 2017.

Tianyi Lin, Nhat Ho, and Michael Jordan. On efficient optimal transport: An analysis of greedy and accelerated mirror descent algorithms. In International Conference on Machine Learning , pages 3982-3991. PMLR, 2019.

Tianlin Liu, Joan Puigcerver, and Mathieu Blondel. Sparsity-constrained optimal transport. International Conference on Learning Representations , 2022.

Dirk A Lorenz, Paul Manns, and Christian Meyer. Quadratically regularized optimal transport. Applied Mathematics & Optimization , 83(3):1919-1949, 2021.

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.

Nicolas Papadakis, Gabriel Peyré, and Edouard Oudet. Optimal transport with proximal splitting. SIAM Journal on Imaging Sciences , 7(1):212-238, 2014.

Hanyu Peng, Mingming Sun, and Ping Li. Optimal transport for long-tailed recognition with learnable cost matrix. 2022.

Gabriel Peyré, Marco Cuturi, et al. Computational optimal transport: With applications to data science. Foundations and Trends® in Machine Learning , 11(5-6):355-607, 2019.

Clarice Poon and Jingwei Liang. Trajectory of alternating direction method of multipliers and adaptive acceleration. Advances in Neural Information Processing Systems , 32, 2019.

Lucas Roberts, Leo Razoumov, Lin Su, and Yuyang Wang. Gini-regularized optimal transport with an application to spatio-temporal forecasting. arXiv preprint arXiv:1712.02512 , 2017.

R Tyrrell Rockafellar and Roger J-B Wets. Variational analysis , volume 317. Springer Science &Business Media, 2009.

Tim Salimans and Durk P Kingma. Weight normalization: A simple reparameterization to accelerate training of deep neural networks. Advances in neural information processing systems , 29, 2016.

Tim Salimans, Han Zhang, Alec Radford, and Dimitris Metaxas. Improving GANs using optimal transport. International Conference on Learning Representations , 2018.

Bernhard Schmitzer. Stabilized sparse scaling algorithms for entropy regularized transport problems. SIAM Journal on Scientific Computing , 41(3):A1443-A1481, 2019.

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.

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.

Richard Sinkhorn and Paul Knopp. Concerning nonnegative matrices and doubly stochastic matrices. Pacific Journal of Mathematics , 21(2):343-348, 1967.

James Thornton and Marco Cuturi. Rethinking initialization of the Sinkhorn algorithm. In International Conference on Artificial Intelligence and Statistics , pages 8682-8698. PMLR, 2023.

Sinong Wang and Ness Shroff. A new alternating direction method for linear programming. Advances in Neural Information Processing Systems , 30, 2017.

Yao-Liang Yu. On decomposing the proximal map. Advances in neural information processing systems , 26, 2013.

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 X ij = 0 implies that (prox ρh ( X )) ij = 0 . Moreover, if X ij > 0 , then (prox ρh ( X )) ij ≥ 0 .

Proof. Let Y = prox ρh ( X ) . By the definition of the proximal operator, every Z ∈ R m × n satisfies

or

Moreover, since h is sparsity promoting h ( Z ) ≤ h ( Y ) , and it holds that

Let Z be equal to Y in all entries but the ( i, j ) th, which is set to zero. If X ij = 0 , then

and invoking (12) gives that Y ij = 0 . If X ij > 0 , on the other hand, then

and (12) gives

which implies that Y ij ≥ 0 . The proof is complete.

Lemma A.1 and Theorem 1 in [42] allows to prove the composition lemma quite easily. We end this section with the proof.

Proof. (Lemma 3.1) For p ( X ) = ι R m × n + ( X ) , it holds that

for all X ≥ 0 , where N R m × n + ( X ) = { Y ∈ R m × n : Y ij = 0 if X ij > 0 , Y ij ≤ 0 if X ij = 0 } .

By Lemma A.1, if X ij = 0 , then (prox ρh ( X )) ij = 0 , and hence ( ∂p ( X )) ij = ( ∂p (prox ρh ( X ))) ij . Further, for positive entries X ij > 0 , ( ∂p ( X )) ij = { 0 } , which implies that ( ∂p (prox ρh ( X ))) ij ⊇ ( ∂p ( X )) ij . Altogether, this means that

and

Theorem 1 in [42] then gives

and

A.2 On the non-degeneracy condition of Asssumption 1

Assumption 1 is violated if Y ∗ is a fixed point to equation (3) and X ⋆ = prox ρf ( Y ⋆ ) , but the regularity condition not fulfilled. That is, if

We can write these conditions more conveniently as

Let us now explore to what degree Assumption 1 can be seen as restrictive when it is invoked in Theorem 1. To this end, we also assume that h is twice continuously differentiable for X > 0 .

Recall that f = ⟨ C, X ⟩ + ι R m × n + + 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 ∂f is set-valued, i.e. when X ∗ ij = 0 for some ( i, j ) . Since h ( X ) is differentiable for X > 0 , we only need to consider the zero entries. If X ∗ ij = 0 , ∂ι R m × n + = N R m × n + implies that ( ∂ι R m × n + \ relint ∂ι R m × n + ) ij = 0 and thus that ∂ι R m × n + \ relint ∂ι R m × n + = 0 . Consequentially, ∂f ( X ⋆ ) \ relint ∂f ( X ⋆ ) = C + ∂h ( X ∗ ) \ relint ∂h ( X ∗ ) , and

By the convergence of the algorithm, equation (5) in the main document implies that there exist two vectors ϕ ⋆ and φ ⋆ such that Y ⋆ -X ⋆ = ϕ ⋆ 1 ⊤ n + 1 m φ ⋆ ⊤ . This gives that the cost is on the form

In other words, it is only cost matrices with a specific (and rather particular) structure that violate Assumption 1. For instance, if h is smooth, e.g. h = λ ∥ · ∥ 2 F , Assumption 1 is only violated when the cost can be written as a sum of two given rank-1 matrices. This is rarely met in real applications.

A.3 Proof of Theorem 1 and 2

To prove the theorem we need the notion of partial smoothness:

Definition A.1 (Partial Smoothness) . A lower semi-continuous function f : R n → R is said to be partly smooth at x ⋆ ∈ R n with respect to the C 2 -manifold M⊂ R n if:

(i) The restriction of f to M : f | M is C 2 around x ⋆ ,

(ii) T M ( x ⋆ ) = (par ∂f ( x ⋆ )) ⊥ , where (par ∂f ( x ⋆ )) ⊥ denotes the orthogonal complement of the smallest subspace parallel to ∂f ( x ⋆ ) ,

(iii) ∂f ( x ) is continuous at x ⋆ relative to M .

Let X ⋆ be a solution to the OT problem that fulfills X k → X ⋆ , and let

twice continuously differentiable over R m × n ++ , by the conditions in the theorem, f is partly smooth with respect to M 1 . Further, since g is constant over M 2 , it is also partly smooth over M 2 . The optimality conditions of the regularized OT problem are

Moreover, since ∂g = N X , and X is an affine subspace, ∂g = relint ∂g . Therefore, Assumption 1 extends to the following the non-degeneracy condition:

With the partial smoothness and the non-degeneracy condition established, we have all the building blocks needed for proving Theorem 1 and 2.

Proof. ( Theorem 1 and 2 ) By invoking the first order optimality condition associated with the first update in the DR-splitting scheme in (4), we get:

or

Similarly, the second update in (4) is associated with the inclusion:

or equivalently

where we used the third equation of the DR-update, i.e Y k +1 -Y k = Z k +1 -X k +1 . Notice how (17) and (18) relate closely to the optimality conditions (15). Now consider the function F ( X,Z ) = f ( X ) + g ( Z ) -ρ -1 ⟨ X -Z, Y ∗ -X ∗ ⟩ . Since X k +1 ∈ R m × n + and Z k +1 ∈ X for all k ≥ 1 , it holds that F ( X k +1 , Z k +1 ) = ⟨ C, X k +1 ⟩ + h ( X k +1 ) -ρ -1 ⟨ X k +1 -Z k +1 , Y ∗ -X ∗ ⟩ . Since f + g are convex, closed, and attains a minimizer, Theorem 25.6 in [3] can be used to establish that Y k → Y ⋆ , when k →∞ , where Y ⋆ is a fixed point of the DR-update. In particular,

which can be used to deduce that the third term of F tends to zero as k →∞ :

Consequentially,

As h is assumed to be convex, proper, and lower semicontinuous, we have that h is subdifferentially continuous in its domain (including at X ⋆ ) [33]. This gives

Here, the first inequality follows from that dist { x, S } = inf {∥ x -y ∥ : y ∈ S } ≤ ∥ x -y ∥ , for all y ∈ S . The second is the Cauchy Schwarz inequality, and the third inequality follows from the non-expansiveness of proximal operators:

Similarly,

Hence, dist { 0 , ∂F ( X k +1 , Z k +1 ) } → 0 . This together with (16) and (19), and the fact that f is partially smooth at X ⋆ with respect to M 1 , we apply Theorem 5.3 in [20] which proves the theorem. To prove Theorem 2, one can directly invoke Theorem 5.6 in [23].

B Initialization

When using the initialization X 0 = pq ⊤ , ϕ 0 = 0 m , and φ 0 = 0 n , and the default stepsize, the updates in (6) gives that

when k = 1 , 2 , . . . N , where

When the cost is normalized, we can use the rough approximation, C ij ∼ 1 , p i ∼ m -1 , and q j ∼ n -1 , to get the simple initialization strategy: X 0 = 0 m × n , ϕ 0 = (3( m + n )) -1 (1 + m/ ( m + n )) 1 m , and φ 0 = (3( m + n )) -1 (1 + n/ ( m + n )) 1 n , that skips these first N iterations.

C Derivation of Dual and Stopping criterion

A simple calculation reveals that g ( X ) = ι X ( X ) has the Fenchel conjugate

Hence, any feasible dual variable is on the form U = µ 1 ⊤ n + 1 m ν ⊤ , for which g ∗ ( U ) = p ⊤ µ + q ⊤ ν . Further, the conjugate of f ( X ) = ⟨ C, X ⟩ + ι R m × n + ( X ) + h ( X ) can be expressed as

Moreover, since h is sparsity promoting,

meaning that f ∗ ( U ) = h ∗ ([ U -C ] + ) . This gives the dual problem

To relate the iterates to the optimality condition, we define X k +1 = prox ρf ( Y k ) , Z k +1 = prox ρg (2 X k +1 -Y k ) , and U k := ( Y k -1 -X k ) /ρ and observe that

By theorem 25.6 in [3], Z k +1 -X k +1 → 0 , and U k +1 → U ⋆ where U ⋆ is a solution to (20). Hence ( ρ -1 ϕ k ) 1 n + 1 m ( ρ -1 φ k ) ⊤ → U ∗ , meaning that ρ -1 ϕ k and ρ -1 φ k will correspond to µ and ν in (20) respectively. Using that r k and s k in (6) are directly related to the primal residual, we can derive the following stopping criteria:

where ∆ h denotes the duality gap deviation compared to the unregularized problem. Further, notice that by letting ¯ X k +1 = [ X k + ϕ k 1 ⊤ n + 1 m φ ⊤ k -ρC ] + , we can express [ ϕ k 1 ⊤ n + 1 m φ ⊤ k -ρC ] + = [ ¯ X k +1 -X k ] + , which facilitates evaluating the duality gap. Notice that h ∗ may not be finite. In 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.

D Backpropagation through regularized OT-costs

We let h c = h + ι R m × n + + ι X , and define the OT-cost as follows: OT h ( C ) = inf X ⟨ C, X ⟩ + h c ( X ) . Note that OT cost is equal to the optimal value of the OT problem (1), i.e. if X ⋆ is a solution to (1), then OT h ( C ) = ⟨ C, X ⋆ ⟩ + h c ( X ⋆ ) . The OT cost is also closely related to the Fenchel conjugate of h c , as OT h ( C ) = -¯ h ∗ ( -C ) , and hence ¯ h ∗ ( -C ) + ¯ h ( X ⋆ ) = ⟨-C, X ⋆ ⟩ . Invoking subdifferential properties of the Fenchel conjugate gives that X ⋆ ∈ ∂ ¯ h ⋆ ( -C ) . Further

In particular, whenever the solution is unique, such as when h is strongly convex, then ∇ C OT h ( C ) = X ⋆ . Therefore, using the transportation plan estimated via RDROT is a 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 basic 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.

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, a 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:

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.

A block is a group of threads that execute together on an SM in the GPU. Threads within a block can cooperate and communicate using shared memory.

The grid refers to the entire set of blocks that will be 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 are scheduled and executed together on a single SM. All threads within a 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 on the first three memory types, since they are the most relevant to the kernel design in the next section.

On NVIDIA GPUs, the global memory refers to the main memory available on the GPU. It is a 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 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 E.2.1.

The shared memory is a low-latency memory space, physically located on the GPU chip. The shared memory is used for data sharing within a 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 characteristics are that it 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.

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 a warp access memory in a 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 a 32-byte, 64-byte, or 128-byte transaction that is aligned with their sizes, respectively. If the threads in a 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), it only costs one transaction. If the continuous memory block is misaligned with the 128-byte pattern, it costs two transactions. If the memory accessed

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 2 . 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-lasso 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 is designed in 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 single-precision floats 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 5 where the old working area of a block from [27] is marked with 'Planned Block Work Area' and the updated working area of a block is marked with 'Actual Block Work Area'. The comparison of the per-iteration runtime before and after the optimization in Table 3 shows that the update solves the problem of uncoalesced memory access, and that the implementation performs more predictably on various problem sizes.

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 to be loaded once for every update. As the update of X is done column by column, the auxillary variables ϕ and φ can be shared and reused among the threads in a block. To speedup the memory access, we load the required elements of ϕ and φ into shared memory at the beginning of the main kernel.

Figure 5 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 of the following steps

Load the corresponding elements of ϕ and φ from the shared memory.

Load the corresponding elements of X and C from the global memory. Compute [ Y ij -ρC ij ] + and C ij · X ij . For the quadratically regularized case, we can directly update X ij with Y ij after multiplying it with the scale 1 / (1 + ρα ) .

Perform in-warp reduction to compute the column sum and the objective value in the block.

Use atomic operations to accumulate the results to the corresponding row and column sums.

This step is only used by the group-lasso kernel, introduced in Section E.2.3.

Similar to the implementation by [27], the iterations are divided into even iterations and odd iterations.

For the even iterations, the kernel updates X k +1 as

For the odd iterations, it updates X k +1 as

In this way, the kernel only needs to 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 a scalar. This minor change introduces negligible changes in 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 5. More specifically, it performs the following steps:

Load the corresponding elements of ϕ and φ from the shared memory.

Load the corresponding elements of X and C from the global memory. Compute [ Y ij -ρC ij ] + and C ij · X ij . Compared with the quadratically regularized DROT, the kernel cannot update X ij until the scale of the group is ready in step (5) .

Perform in-warp reduction. For the group-lasso regularized case, we also reduce the sum of [ Y ij -ρC ij ] 2 + to compute the norm of the group.

Accumulate the results to the corresponding sum of the row, the sum of the column, and the sum of the squared elements within the group using atomic operations.

Compute the scale for the group and apply it to the computed [ Y ij -ρC ij ] + so as to update X ij . (For large problem sizes, where a thread block can only access part of the group, the process of reduction and broadcast is done in an additional kernel function).

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 not fit in the shared memory, it is inevitable to access the matrix C in every iteration. Both factors lead to an increase in the per-iteration 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-Knopp 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-Knopp algorithm, 12 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 2 and 4 classes, respectively. Note that the Sinkhorn-based implementation of group-lasso OT in POT [16] does not return the number of iterations, nor reports the per-iteration run-time. This makes

Cost

C

Transport Plan

X

(

2

)

(

3

)

(

4

)

(

4

)

ϕ

(

1

)

φ

(

1

)

Sums of Cols

X

T

j

Sums of Rows

X

i

Sum of Group Norms

G

j,k

(

5

)

Planned Block Work Area

Actual Block Work Area

Loaded in Shared Memory

Update

X

with 16 Threads

Within-Warp Reduction

it impossible to 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 and ours, we cover four variants of the Sinkhorn-Knopp algorithm in the comparison. The implementation of Sinkhorn-Knopp 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 SinkhornKnopp algorithm is improved by moving the computation to log scale. The decision to 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-Knopp algorithm that we evaluate. Since our methods work directly on X we let them compute the objective value and the primal residual in every iteration.

The results show that our QDROT implementation has roughly the same per-iteration runtime as the standard Sinkhorn-Knopp algorithm. The per-iteration runtime of the Sinkhorn-Knopp algorithm in log scale is roughly 4.3-9.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 regularization on more datasets of various sizes and using different hyperparameters. The results are

Sinkhorn

per iteration (ms)

Sinkhorn w/o stopping crit.

8

QDROT

Sinkhorn (log)

Sinkhorn (log)w/o stopping crit.

6

GLDROT (2 classes)

GLDROT (4 classes)

4

Runtime

2

0

2000

4000

6000

8000

10000

Problem size (m = n)

consistent with the experiments that could be included in the main document and demonstrate significant advantages of our method compared to the state-of-the-art.

Quadratic regularization In the main manuscript, we compared our solver with an L-BFGS method applied to the dual problem. Two problem sizes were considered: m = 1000 , n = 1000 , and m = 2000 , n = 3000 . In Figure 2, we rerun the benchmark on 4 additional problem sizes. Notice that all results are consistent with those presented in the paper, showing that our results generalize to more settings beyond the ones presented in the paper.

(a)

m

= 100

,

n

= 100

(b)

m

= 1000

,

n

= 1000

(c)

m

= 2000

,

n

= 3000

(d)

m

= 10000

,

n

= 10000

10

QDROT

LBFGS

LBFGS

QDROT

10

Run time (ms)

Run time (ms)

102

101

100

100

Regularization

Regularization

LBFGS

QDROT

105

+

HH

QDROT

LBFGS

Run time (ms)

Run time (ms)

10°

10²

103

白

白

+

102

101

01

Regularization

Regularization

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 m = 1000 , n = 1500 , for a 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, GLDROT consistently outperforms the state-of-the-art in all considered problems.

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 the MNIST and CIFAR10 datasets.

To define the loss, we need the following sample-based OT-cost.

Here m and n are the batch sizes corresponding to X and Y , and C X , Y is a matrix with pairwise distances between the samples of X and Y . The mini-batch energy distance [35] is then given by

Here, X , X ′ are two independent mini-batches from real data, while Y , Y ′ are two independent mini-batches generated from the generator.

In the experiments, we used the cosine similarity to parameterize the cost, and we replaced the OT solver used in [35] - the Sinkhorn-Knopp algorithm, by the PyTorch wrapper for our OT solver with the quadratic regularizer.

The model structures are adapted from the ones in [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 3 × 10 -4 , with β 1 = 0 . 5 and β 2 = 0 . 999 . The batch size is 1024 for the MNIST experiment and 2048 for the CIFAR10 experiment. We update the generator three times for every discriminator update. As for the parameters of our OT solver, we use ϵ = 10 -4 as 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 are not the focus of this paper, we leave such refinements as future work.

G.1 MNIST Architecture and Results

B

S

7

力

力

1

7

九

7

5

2

6

9

2

3

S

?

3

9

3

6

The generated samples in Figure 8 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.

G.2 CIFAR10 Architecture and Results

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 Sinkhorn-Knopp algorithm as the OT solver in the mini-batch energy loss.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nThe 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.\n\nQUESTIONS:\nl.143: \"since $Y_k$ can be eliminated...\": could you elaborate on what you mean? After multiple readings, this remains unclear.\n\nLIMITATIONS:\nNo negative societal impact for this work.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-11"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nThe 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\n\nQUESTIONS:\nIf 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.\n\nLIMITATIONS:\nA discussion on the tradeoffs between the proposed algorithm and minimizing sinkhorn divergences would certainly help to correctly evaluate the proposed contribution.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-06-29"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nTo 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.\n\nWEAKNESSES:\nThe 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\nQUESTIONS:\nMy 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\nLIMITATIONS:\nNA", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn 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).\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nThe 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.\n\nQUESTIONS:\nWhat happens if we use the entropic regularizer in RDROT?\n\nLIMITATIONS:\nSee weaknesses.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-13"}