SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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\|_{\infty} = 1$ , a stepsize that performs well regardless of  $h$  is  $\rho = 2(m + n)^{-1}$ . This stepsize was proposed for DR-splitting 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^\top$ , and  $\phi_0 = \mathbf{0}_m$ ,  $\varphi_0 = \mathbf{0}_n$ . However, when using the stepsize specified above, this will result in  $X_k = \mathbf{0}_{m \times n}$  for  $1 \leq k \leq N$ , where  $N = O(\max(m, n))$ . By initializing  $\phi_0 = (3(m + n))^{-1}(1 + m / (m + n))\mathbf{1}_m$  and  $\varphi_0 = (3(m + n))^{-1}(1 + n / (m + n))\mathbf{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  $\rho^{-1}\phi_k$ ,  $\rho^{-1}\varphi_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 suboptimality bounds are costly to compute from the generated iterates. This is usually addressed by only computing the primal residuals  $\| X_{k}\mathbf{1}_{m} - p\|$ , and  $\| X_k^\top \mathbf{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  $\phi_{k}$  and  $\varphi_{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<sup>1</sup>. 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 - \rho 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  $\bar{X}_{k+1,g}$ ,  $g \in \mathcal{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  $\bar{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^{\star}$  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  $\rho^{-1}(\phi^{\star};\varphi^{\star})$  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}^{\star}$  and  $\rho^{-1}(\phi^{\star};\varphi^{\star})$  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.
All experiments mentioned below are performed using a single NVIDIA Tesla V100 GPU.
The implementation is publicly available at GitHub: https://github.com/WangZesen/Regularized-DROT
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 \times 500$  and  $5000 \times 5000$ , respectively. We generated cost matrices with  $C_{ij} = \frac{1}{2}\|\mathbf{x}_i - \mathbf{x}_j\|_2^2$ , where  $\mathbf{x}_i, \mathbf{x}_j$  are simulated 2D Gaussian variables with random parameters specific for each data set. We used the quadratic regularization term  $((m + n)\alpha / 2)\|X\|_F^2$  with  $\alpha$  ranging between 0.0005 to 0.2. For QDROT, we included  $\alpha = 0$  for reference. We ran the algorithms for each dataset and regularization parameter until they reached an accuracy level of 0.0001. For completeness, 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$
Figure 2: Comparison between the RDROT (QDROT) using the GPU kernel, a PyTorch version, and an L-BFGS method applied to the dual for different quadratic regularization parameters run on GPU. 50 datasets of two different sizes were simulated.
Group Lasso and Domain Adaptation
OT is the backbone of many domain adaption techniques since it provides a natural framework to estimate correspondences between a source domain and a target domain. These correspondences can subsequently be used to align the respective domains to improve the test accuracy. Specifically, for a given source data set  $\{(\mathbf{x}_s^i,y^i)\}_{i = 1}^m$  and a target set  $\{\mathbf{x}_t^j\}_{j = 1}^n$ , where  $\{y^{i}\}_{i = 1}^{m}$  are data labels, we define the cost matrix elements  $C_{ij} = d(\mathbf{x}_s^i,\mathbf{x}_t^j)$  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  $\mathbf{x}_{s,a}^{i} = \frac{1}{p_{i}}\sum_{j = 1}^{m}X_{ij}\mathbf{x}_{t}^{j}$ . 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 P0T 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} = \frac{1}{2}\| \mathbf{x}_s^i -\mathbf{x}_t^j\|^2$  was used and normalized so that  $\| C\|_{\infty} = 1$ . To compare the performance of
Figure 3: A toy example illustrating the effect strongly convex regularizations have on the support of the resulting adapted domains. Pink and purple points mark the labeled training set and the gray points the test domain. Notice how the Sinkhorn-based method, as well as the Quadratic Regularized method transfers the source domain samples towards the center of the target domain, while our approach transfers the samples in a way that matches the true variation of the test data
the algorithms, we computed the time taken to reach a tolerance of  $10^{-4}$  and evaluated 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
Table 1: Performance benchmark of RDROT for Group-Lasso regularization (GLDROT) against the conditional-gradient based Sinkhorn algorithm. Median and 10th and 90th percentile statistics are included. Our method is competitive in terms of both speed and adaption quality.
Method	Reg.	Runtime (s) ↓	Agg. W2 dist. ↓
Ent	GL	Median	q10	q90	Median	q10	q90
GLSK [10]	1e-3	1e-3	3.77	3.68	8.90	0.311	0.0657	6.48
GLSK [10]	1e-1	1e-3	3.73	3.71	3.77	8.24	3.47	31.6
GLSK [10]	1e+2	1e-3	1.86	1.86	1.88	52.8	10.9	311
GLDROT (ours)	-	1e-3	0.0619	0.0475	0.0771	0.0576	0.0262	0.182
GLDROT (ours)	-	5e-3	0.0384	0.0306	0.0474	0.0879	0.0358	0.274
above, we compare the per-iteration costs for several OT solvers in the supplementary material. These strengthen our claim further that our algorithm achieves "lightspeed" even for large problem sizes.
Training of generative models
As an illustration of the efficiency of the GPU version of the algorithm, we use our PyTorch wrapper to implement the Minibatch Energy Distance [35], 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.
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,
Figure 4: Generated samples by GANs trained on MNIST and CIFAR10 datasets with the Minibatch Energy Distance. The original sizes of samples are (28,28) and (32,32), respectively.
and whose operations are readily parallelizable on GPUs. The performance of our algorithm is supported by both theoretical arguments and extensive empirical evaluations. In particular, the experiments demonstrate that the resulting algorithms can be up to two orders of magnitude faster than the current state-of-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.
This work was supported in part by the Knut and Alice Wallenberg Foundation, the Swedish Research Council, and the Wallenberg AI, Autonomous Systems and Software Program (WASP). The computations were enabled by resources provided by the Swedish National Infrastructure for Computing (SNIC), partially funded by the Swedish Research Council through grant agreement no. 2018-05973.
[36] Bernhard Schmitzer. Stabilized sparse scaling algorithms for entropy regularized transport problems. SIAM Journal on Scientific Computing, 41(3):A1443-A1481, 2019.
[37] Vivien Seguy, Bharath Bhushan Damodaran, Rémi Flamary, Nicolas Courty, Antoine Rolet, and Mathieu Blondel. Large-scale optimal transport and mapping estimation. International Conference on Learning Representations, 2018.
[38] Wenling Shang, Kihyuk Sohn, Diogo Almeida, and Honglak Lee. Understanding and improving convolutional neural networks via concatenated rectified linear units. In international conference on machine learning, pages 2217-2225. PMLR, 2016.
[39] Richard Sinkhorn and Paul Knopp. Concerning nonnegative matrices and doubly stochastic matrices. Pacific Journal of Mathematics, 21(2):343-348, 1967.
[40] James Thornton and Marco Cuturi. Rethinking initialization of the Sinkhorn algorithm. In International Conference on Artificial Intelligence and Statistics, pages 8682-8698. PMLR, 2023.
[41] Sinong Wang and Ness Shroff. A new alternating direction method for linear programming. Advances in Neural Information Processing Systems, 30, 2017.
[42] Yao-Liang Yu. On decomposing the proximal map. Advances in neural information processing systems, 26, 2013.
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  $\rho > 0$ , then  $X_{ij} = 0$  implies that  $(\operatorname{prox}_{\rho h}(X))_{ij} = 0$ . Moreover, if  $X_{ij} > 0$ , then  $(\operatorname{prox}_{\rho h}(X))_{ij} \geq 0$ .
Proof. Let  $Y = \operatorname{prox}_{\rho h}(X)$ . By the definition of the proximal operator, every  $Z \in \mathbb{R}^{m \times n}$  satisfies
$$
h (Y) + \frac {1}{2 \rho} \| Y - X \| _ {F} ^ {2} \leq h (Z) + \frac {1}{2 \rho} \| Z - X \| _ {F} ^ {2}
$$
$$
h (Y) + \frac {1}{2 \rho} \left(\| Y - X \| _ {F} ^ {2} - \| Z - X \| _ {F} ^ {2}\right) \leq h (Z).
$$
Moreover, since  $h$  is sparsity promoting  $h(Z) \leq h(Y)$ , and it holds that
$$
\frac {1}{2 \rho} \left(\| Y - X \| _ {F} ^ {2} - \| Z - X \| _ {F} ^ {2}\right) \leq 0. \tag {12}
$$
Let  $Z$  be equal to  $Y$  in all entries but the  $(i,j)$ th, which is set to zero. If  $X_{ij} = 0$ , then
$$
\frac {1}{2 \rho} \left(\| Y - X \| _ {F} ^ {2} - \| Z - X \| _ {F} ^ {2}\right) = \frac {1}{2 \rho} Y _ {i j} ^ {2}
$$
and invoking (12) gives that  $Y_{ij} = 0$ . If  $X_{ij} > 0$ , on the other hand, then
$$
\frac {1}{2 \rho} \bigg (\| Y - X \| _ {F} ^ {2} - \| Z - X \| _ {F} ^ {2} \bigg) = \frac {1}{2 \rho} \bigg ((Y _ {i j} - X _ {i j}) ^ {2} - X _ {i j} ^ {2} \bigg),
$$
and (12) gives
$$
\left(Y _ {i j} - X _ {i j}\right) ^ {2} \leq X _ {i j} ^ {2}
$$
which implies that  $Y_{ij} \geq 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) = \iota_{\mathbb{R}_+^{m\times n}}(X)$ , it holds that
$$
\operatorname {p r o x} _ {\rho p} (X) = [ X ] _ {+}
$$
$$
\partial p (X) = N _ {\mathbb {R} ^ {n} \times n} (X) \supseteq \{0 \}
$$
for all  $X \geq 0$ , where  $N_{\mathbb{R}_+^{m \times n}}(X) = \{Y \in \mathbb{R}^{m \times n} : Y_{ij} = 0 \text{ if } X_{ij} > 0, Y_{ij} \leq 0 \text{ if } X_{ij} = 0\}$ .
By Lemma A.1, if  $X_{ij} = 0$ , then  $(\mathrm{prox}_{\rho h}(X))_{ij} = 0$ , and hence  $(\partial p(X))_{ij} = (\partial p(\mathrm{prox}_{\rho h}(X)))_{ij}$ . Further, for positive entries  $X_{ij} > 0$ ,  $(\partial p(X))_{ij} = \{0\}$ , which implies that  $(\partial p(\mathrm{prox}_{\rho h}(X)))_{ij} \supseteq (\partial p(X))_{ij}$ . Altogether, this means that
$$
\partial p \left(\operatorname {p r o x} _ {\rho h} (X)\right) \supseteq \partial p (X).
$$
Theorem 1 in [42] then gives
$$
\operatorname {p r o x} _ {\rho p + \rho h} (X) = \operatorname {p r o x} _ {\rho h} \left(\operatorname {p r o x} _ {\rho p} (X)\right) = \operatorname {p r o x} _ {\rho h} ([ X ] _ {+})
$$
$$
\begin{array}{l} \operatorname {p r o x} _ {\rho f} (X) = \operatorname {p r o x} _ {\rho \langle C, \cdot \rangle + \rho p + \rho h} (X) \\ = \operatorname {p r o x} _ {\rho p + \rho h} (X - \rho C) \\ = \operatorname {p r o x} _ {\rho h} \left(\left[ X - \rho C \right] _ {+}\right). \\ \end{array}
$$
A.2 On the non-degeneracy condition of Assumption 1
Assumption 1 is violated if  $Y^*$  is a fixed point to equation (3) and  $X^{\star} = \mathrm{prox}_{\rho f}(Y^{\star})$ , but the regularity condition not fulfilled. That is, if
$$
\frac {1}{\rho} \left(Y ^ {*} - X ^ {*}\right) \in \partial f (X ^ {\star}), \text {w h i l e} \frac {1}{\rho} \left(Y ^ {*} - X ^ {*}\right) \notin \operatorname {r e l i n t} \partial f (X ^ {\star}).
$$
We can write these conditions more conveniently as
$$
\frac {1}{\rho} \left(Y ^ {*} - X ^ {*}\right) \in \partial f \left(X ^ {\star}\right) \backslash \operatorname {r e l i n t} \partial f \left(X ^ {\star}\right). \tag {13}
$$
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 = \langle C, X \rangle + \iota_{\mathbb{R}_+^{m \times n}} + h(X)$  and note that the inclusion (13) cannot be satisfied if  $X^* > 0$ , since  $\partial f(X^\star)$  is a singleton. We therefore focus on the case when  $\partial 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$ ,  $\partial \iota_{\mathbb{R}_+^{m \times n}} = N_{\mathbb{R}_+^{m \times n}}$  implies that  $(\partial \iota_{\mathbb{R}_+^{m \times n}} \setminus \operatorname{relint} \partial \iota_{\mathbb{R}_+^{m \times n}})_{ij} = 0$  and thus that  $\partial \iota_{\mathbb{R}_+^{m \times n}} \setminus \operatorname{relint} \partial \iota_{\mathbb{R}_+^{m \times n}} = 0$ . Consequently,  $\partial f(X^\star) \setminus \operatorname{relint} \partial f(X^\star) = C + \partial h(X^*) \setminus \operatorname{relint} \partial h(X^*)$ , and
$$
\frac {1}{\rho} \left(Y ^ {*} - X ^ {*}\right) = C + g, \text {w h e r e} g \in \partial h (X ^ {*}) \backslash \operatorname {r e l i n t} \partial h (X ^ {*}).
$$
By the convergence of the algorithm, equation (5) in the main document implies that there exist two vectors  $\phi^{\star}$  and  $\varphi^{\star}$  such that  $Y^{\star} - X^{\star} = \phi^{\star}\mathbf{1}_{n}^{\top} + \mathbf{1}_{m}\varphi^{\star \top}$ . This gives that the cost is on the form
$$
C = \rho^ {- 1} \left(\phi^ {\star} \mathbf {1} _ {n} ^ {\top} + \mathbf {1} _ {m} \varphi^ {\star^ {\top}}\right) - g.
$$
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 = \lambda \| \cdot \|_F^2$ , 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: \mathbb{R}^n \to \mathbb{R}$  is said to be partly smooth at  $x^{\star} \in \mathbb{R}^{n}$  with respect to the  $C^2$ -manifold  $\mathcal{M} \subset \mathbb{R}^n$  if:
(i) The restriction of  $f$  to  $\mathcal{M}$ :  $f|_{\mathcal{M}}$  is  $C^2$  around  $x^\star$ ,
(ii)  $\mathcal{T}_{\mathcal{M}}(x^{\star}) = (\operatorname {par}\partial f(x^{\star}))^{\perp}$ , where  $(\operatorname {par}\partial f(x^{\star}))^{\perp}$  denotes the orthogonal complement of the smallest subspace parallel to  $\partial f(x^{\star})$
(iii)  $\partial f(x)$  is continuous at  $x^*$  relative to  $\mathcal{M}$ .
Let  $X^{\star}$  be a solution to the OT problem that fulfills  $X_{k} \to X^{\star}$ , and let
$$
\mathcal {M} _ {1} = \left\{X \in \mathbb {R} _ {+} ^ {m \times n}: X _ {i j} = 0, \text {i f} X _ {i j} ^ {\star} = 0, X _ {i j} > 0 \text {i f} X _ {i j} ^ {\star} > 0 \right\}, \quad \mathcal {M} _ {2} = \mathcal {X}. \tag {14}
$$
Recall that  $f(X) = \langle C, X \rangle + \iota_{\mathbb{R}_+^{m \times n}}(X) + h(X)$ , and  $g(X) = \iota_{\mathcal{X}}(X)$ . Since  $h$  is assumed to be twice continuously differentiable over  $\mathbb{R}_{++}^{m \times n}$ , by the conditions in the theorem,  $f$  is partly smooth with respect to  $\mathcal{M}_1$ . Further, since  $g$  is constant over  $\mathcal{M}_2$ , it is also partly smooth over  $\mathcal{M}_2$ . The optimality conditions of the regularized OT problem are
$$
\begin{array}{l} 0 \in \partial f (X ^ {\star}) - \frac {1}{\rho} \left(Y ^ {*} - X ^ {*}\right) \tag {15} \\ 0 \in \partial g (X ^ {\star}) + \frac {1}{\rho} \left(Y ^ {*} - X ^ {*}\right). \\ \end{array}
$$