SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Jacob Lindbäck  
EECS, KTH
Stockholm, Sweden jlindbac@kth.se
Zesen Wang  
EECS, KTH
Stockholm, Sweden  
zesen@kth.se
Mikael Johansson  
EECS, KTH
Stockholm, Sweden  
mikaelj@kth.se
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.
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  $100\mathrm{x}$  speed-up for several applications.
For discrete probability distributions, solving regularized OT problems amounts to finding a solution to the optimization problem
$$
\text {m i n i m i z e} \quad \langle C, X \rangle + h (X)
$$
$$
\begin{array}{l} X \in \mathbb {R} _ {+} ^ {m \times n} \\ \text {s u b j e c t t o} \quad X \mathbf {1} _ {n} = p, X ^ {\top} \mathbf {1} _ {m} = q. \end{array} \tag {1}
$$
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) = \epsilon \sum_{ij} X_{ij} \log X_{ij}$ , where  $\epsilon > 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  $\epsilon$  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  $\epsilon$  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.
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.
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 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 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.
For any  $X,Y\in \mathbb{R}^{m\times n}$ , we let  $\langle X,Y\rangle \coloneqq \mathrm{tr}(X^{\top}Y)$  and  $\| X\| _F\coloneqq \sqrt{\langle X,X\rangle}$ .  $\mathbb{R}_+^{m\times n}$  and  $\mathbb{R}_{-}^{m\times n}$  denote the set of  $m\times n$  matrices with non-negative entries and non-positive entries respectively. We let  $\iota_{S}$  denote the indicator function over a closed convex set  $S$ , i.e  $\iota_S(X) = 0$  if  $X\in S$  and  $\iota_S(X) = \infty$  if  $X\notin S$ , and the relative interior of a set  $S$  is denoted relint  $S$ . The projection onto  $\mathbb{R}_+^{m\times n}$  is denoted  $[\cdot ]_+$ , which sets all negative entries of the inputted matrix to zero. The subdifferential of an extended-real valued function  $h$  is denoted  $\partial h(X)$ , and its proximal operator is defined  $\operatorname {prox}_{\rho h}(x) = \operatorname *{argmin}_z h(z) + \frac{1}{2\rho}\| z - x\|^2$ , where  $\rho >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: \mathbb{R}^{m \times n} \to \mathbb{R} \cup \{+\infty\}$  is said to be sparsity promoting if, for any  $X \in \mathbb{R}^{m \times n}$ ,  $h(X) \geq h(X_s)$  for every  $X_s \in \mathbb{R}^{m \times n}$  with  $(X_s)_{ij} \in \{0, X_{ij}\}$ .
Notice that this function class does not necessarily induce sparsity. For instance  $h = 0$ , or  $h = \| \cdot \|_F^2$  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) = \sum_{g\in \mathcal{G}}\| X_g\| _F$  (group lasso OT)
-  $h(X) = \sum_{ij} X_{ij} \operatorname{arcsinh}(X_{ij} / \beta) - (X_{ij}^2 + \beta^2)^{1/2}$  (hypentropic regularization)
-  $h(X) = \sum_{ij} w_{ij} |X_{ij}|$ , where  $w_{ij} \geq 0$  (weighted  $\ell_1$ -regularization)
$h(X) = \sum_{(ij)\in S}\iota_{X_{ij} = 0}(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  $\| X - A\| _F^2$  can be converted to only involve sparsity promoting regularizers. Note that the negative Shannon entropy  $\sum 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
$$
\underset {x \in \mathbb {R} ^ {n}} {\text {m i n i m i z e}} f (x) + g (x) \tag {2}
$$
using the fixed point update  $y_{k + 1} = T(y_k)$ , where
$$
T (y) = y + \operatorname {p r o x} _ {\rho g} \left(2 \operatorname {p r o x} _ {\rho f} (y) - y\right) - \operatorname {p r o x} _ {\rho f} (y) \tag {3}
$$
and  $\rho > 0$  is a stepsize parameter. If  $y^{\star}$  is a fixed point of  $T$  then  $x^{\star} = \mathrm{prox}_{\rho f}(y^{\star})$  is a solution to (2). Often a third iterate,  $z_{k}$ , is introduced, and the Douglas-Rachford iterations are expressed as
$$
x _ {k + 1} = \operatorname {p r o x} _ {\rho f} \left(y _ {k}\right), \quad z _ {k + 1} = \operatorname {p r o x} _ {\rho g} \left(2 x _ {k + 1} - y _ {k}\right) \quad y _ {k + 1} = y _ {k} + z _ {k + 1} - x _ {k + 1}. \tag {4}
$$
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  $\mathcal{X} = \{X\in \mathbb{R}^{m\times n}:X\mathbf{1}_n = p,X^\top \mathbf{1}_m = q\}$  and introduce the indicator functions  $\iota_{\mathbb{R}_{+}^{m\times n}}$  and  $\iota_{\mathcal{X}}$ , that correspond to the constraints of (1). Recognizing that the projection onto  $\mathcal{X}$  is tractable (by invoking recently derived formulas for matrix projections [4]) the authors of [27] proposed the splitting
$$
f (X) = \langle C, X \rangle + \iota_ {\mathbb {R} _ {+} ^ {m \times n}} (X) \quad \text {a n d} \quad g (X) = \iota_ {\mathcal {X}} (X)
$$
and demonstrated that the update (3) can be simplified to
$$
X _ {k + 1} = \left[ Y _ {k} - \rho C \right] _ {+}, \quad Y _ {k + 1} = X _ {k + 1} + \phi_ {k + 1} \mathbf {1} _ {n} ^ {\top} + \mathbf {1} _ {m} \varphi_ {k + 1} ^ {\top} \tag {5}
$$
where  $\phi_{k}$  and  $\varphi_{k}$  are vectors given via the iterative scheme:
$$
r _ {k + 1} = X _ {k + 1} \mathbf {1} _ {n} - p, \quad s _ {k + 1} = X _ {k + 1} ^ {\top} \mathbf {1} _ {m} - q, \quad \eta_ {k + 1} = f ^ {\top} r _ {k + 1} / (m + n),
$$
$$
\theta_ {k + 1} = \theta_ {k} - \eta_ {k + 1}, \qquad a _ {k + 1} = a _ {k} - r _ {k + 1}, \qquad b _ {k + 1} = b _ {k} - s _ {k + 1},
$$
$$
\phi_ {k + 1} = \left(a _ {k} - 2 r _ {k + 1} + \left(2 \eta_ {k + 1} - \theta_ {k}\right) \mathbf {1} _ {m}\right) / n,
$$
$$
\varphi_ {k + 1} = \left(b _ {k} - 2 s _ {k + 1} + \left(2 \eta_ {k + 1} - \theta_ {k}\right) \mathbf {1} _ {n}\right) / m. \tag {6}
$$
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
$$
X _ {k + 1} = \left[ X _ {k} + \phi_ {k} \mathbf {1} _ {n} ^ {\top} + \mathbf {1} _ {m} \varphi_ {k} ^ {\top} - \rho C \right] _ {+}.
$$
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
$$
f (X) = \langle C, X \rangle + \iota_ {\mathbb {R} _ {+} ^ {m \times n}} (X) + h (X), \text {a n d} g (X) = \iota_ {\chi} (X), \tag {7}
$$
the first update of (4) reads:
$$
X _ {k + 1} = \operatorname {p r o x} _ {\rho f} \left(Y _ {k}\right) = \operatorname {p r o x} _ {\rho h + \iota_ {\mathbb {R} _ {+} ^ {m \times n}}} \left(Y _ {k} - \rho C\right).
$$
This update is efficient given that  $\mathrm{prox}_{\rho h + \iota_{\mathbb{R}_+^m\times n}}(\cdot)$  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) = \langle C, X \rangle + \iota_{\mathbb{R}_+^{m \times n}}(X) + h(X)$  where  $h(X)$  is sparsity promoting, closed, convex and proper when  $X \in \mathbb{R}_+^{m \times n}$ . Then  $\mathrm{prox}_{\rho f}(X) = \mathrm{prox}_{\rho h}([X - \rho 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:
$$
X _ {k + 1} = \operatorname {p r o x} _ {\rho h} \left(\left[ Y _ {k} - \rho C \right] _ {+}\right), \quad Y _ {k + 1} = X _ {k + 1} + \phi_ {k + 1} \mathbf {1} _ {n} ^ {\top} + \mathbf {1} _ {m} \varphi_ {k + 1} ^ {\top}. \tag {8}
$$
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) = \frac{\alpha}{2} \| X\| _F^2$  yields
$$
X _ {k + 1} = (1 + \rho \alpha) ^ {- 1} \left[ Y _ {k} - \rho C \right] _ {+} \tag {9}
$$
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  $\alpha \rightarrow 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) = \lambda \sum_{g\in \mathcal{G}}\| X_g\| _F$ , where  $\mathcal{G}$  is a collection of disjoint index sets, the RDROT update becomes
$$
\bar {X} _ {k + 1} = \left[ Y _ {k} - \rho C \right] _ {+}, \quad X _ {k + 1, g} = \left[ 1 - \frac {\lambda}{\| \bar {X} _ {k + 1 , g} \| _ {F}} \right] _ {+} \bar {X} _ {k + 1, g}, \quad g \in \mathcal {G}. \tag {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 a sparsity-inducing group-lasso regularizer in (1). In this setting,  $\mathcal{G}$  is a collection of sets, each specifying which columns correspond to a particular label. Consequently, 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  $\epsilon$ -accurate solution in  $O(1 / \epsilon)$  iterations for general convex problems [21]. This is a significant improvement over Sinkhorn's iteration complexity of  $O(1 / \epsilon^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.
Figure 1: An illustration of the convergence behavior of the algorithm for a quadratically regularized problem. The stopping criterion and the number of positive entries were computed in every iteration. Notice that the iteration number at 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 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^{\star}$  be a fixed point of (3) such that  $Y_{k}\to Y^{\star}$ , and let  $X^{\star} = \mathrm{prox}_{\rho f}(Y^{\star}) = \mathrm{prox}_{\rho h}([Y^{\star} - \rho C]_{+})$ . We then assume that:
$$
\frac {1}{\rho} \left(Y ^ {*} - X ^ {*}\right) \in \operatorname {r e l i n t} \partial f \left(X ^ {\star}\right). \tag {11}
$$
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 \geq 1$  such that for all  $k \geq K$ ,  $X_{k,ij} = 0$  if and only if  $X_{ij}^{\star} = 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 \geq K$ , then
$$
\left\| X _ {k + 1} - X ^ {\star} \right\| _ {F} \leq r ^ {k - K} \left\| Y _ {K} - Y ^ {\star} \right\| _ {F}, \quad r \in (0, 1).
$$
Moreover, the optimal rate is independent of the stepsize  $\rho$ .
The regularizer  $h$  is locally polyhedral, for instance, in the unregularized setting, and when weighted  $\ell_1$  regularization is used (i.e.  $h(X) = \sum_{ij} w_{ij}|X_{ij}|$ , where  $w_{ij} \geq 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.