SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Moreover, since  $\partial g = N_{\mathcal{X}}$ , and  $\mathcal{X}$  is an affine subspace,  $\partial g = \operatorname{relint} \partial g$ . Therefore, Assumption 1 extends to the following the non-degeneracy condition:
$$
0 \in \operatorname {r e l i n t} \partial f \left(X ^ {\star}\right) - \frac {1}{\rho} \left(Y ^ {*} - X ^ {*}\right) \tag {16}
$$
$$
0 \in \operatorname {r e l i n t} \partial g (X ^ {\star}) + \frac {1}{\rho} \left(Y ^ {*} - X ^ {*}\right).
$$
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:
$$
X _ {k + 1} = \operatorname {p r o x} _ {\rho f} \left(Y _ {k}\right) \Longrightarrow 0 \in \frac {1}{\rho} \left(X _ {k + 1} - Y _ {k}\right) + \partial f \left(X _ {k + 1}\right)
$$
$$
\frac {1}{\rho} \left(Y _ {k} - X _ {k + 1}\right) \in \partial f \left(X _ {k + 1}\right). \tag {17}
$$
Similarly, the second update in (4) is associated with the inclusion:
$$
Z _ {k + 1} = \operatorname {p r o x} _ {\rho g} \left(2 X _ {k + 1} - Y _ {k}\right) \Longrightarrow 0 \in \frac {1}{\rho} \left(Z _ {k + 1} - 2 X _ {k + 1} + Y _ {k}\right) + \partial g \left(Z _ {k + 1}\right)
$$
or equivalently
$$
\frac {1}{\rho} \left(Y _ {k} - X _ {k + 1}\right) - \frac {1}{\rho} \left(Y _ {k + 1} - Y _ {k}\right) \in g \left(Z _ {k + 1}\right). \tag {18}
$$
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) - \rho^{-1}\langle X - Z,Y^{*} - X^{*}\rangle$ . Since  $X_{k + 1}\in \mathbb{R}_{+}^{m\times n}$  and  $Z_{k + 1}\in \mathcal{X}$  for all  $k\geq 1$ , it holds that  $F(X_{k + 1},Z_{k + 1}) = \langle C,X_{k + 1}\rangle +h(X_{k + 1}) - \rho^{-1}\langle X_{k + 1} - Z_{k + 1},Y^{*} - X^{*}\rangle$ . Since  $f + g$  are convex, closed, and attains a minimizer, Theorem 25.6 in [3] can be used to establish that  $Y_{k}\to Y^{\star}$ , when  $k\to \infty$ , where  $Y^{\star}$  is a fixed point of the DR-update. In particular,
$$
\left\| Y _ {k + 1} - Y _ {k} \right\| = \left\|\left(Y _ {k + 1} - Y ^ {\star}\right) - \left(Y _ {k} - Y ^ {\star}\right)\right\| \leq \left\| Y _ {k + 1} - Y ^ {\star} \right\| + \left\| Y _ {k} - Y ^ {\star} \right\|\rightarrow 0,
$$
which can be used to deduce that the third term of  $F$  tends to zero as  $k \to \infty$ :
$$
\left|\left\langle X _ {k + 1} - Z _ {k + 1}, Y ^ {*} - X ^ {*} \right\rangle\right| \leq \| Z _ {k + 1} - X _ {k + 1} \| \| Y ^ {*} - X ^ {*} \| = \| Y _ {k + 1} - Y _ {k} \| \| Y ^ {*} - X ^ {*} \| \rightarrow 0.
$$
$$
F \left(X _ {k + 1}, Z _ {k + 1}\right)\rightarrow \langle C, X ^ {\star} \rangle + h \left(X ^ {\star}\right) = f \left(X ^ {\star}\right) + g \left(X ^ {\star}\right). \tag {19}
$$
As  $h$  is assumed to be convex, proper, and lower semicontinuous, we have that  $h$  is subdifferentially continuous in its domain (including at  $X^{\star}$ ) [33]. This gives
$$
\begin{array}{l} \left. \operatorname {d i s t} \left\{0, \partial_ {X} F \left(X _ {k + 1}, Z _ {k + 1}\right) \right\} = \operatorname {d i s t} \left\{0, \partial f \left(X _ {k + 1}\right) - \frac {1}{\rho} \left(Y ^ {*} - X ^ {*}\right) \right\} \right. \\ \leq \rho^ {- 1} \| \left(Y _ {k} - Y ^ {*}\right) - \left(X _ {k + 1} - X ^ {*}\right) \| \\ \leq \rho^ {- 1} \| Y _ {k} - Y ^ {*} \| + \rho^ {- 1} \| X _ {k + 1} - X ^ {*} \| \\ \leq 2 \rho^ {- 1} \| Y _ {k} - Y ^ {*} \| \\ \rightarrow 0. \\ \end{array}
$$
Here, the first inequality follows from that  $\mathrm{dist}\{x,S\} = \inf \{\| x - y\| :y\in S\} \leq \| x - y\|$ , for all  $y\in S$ . The second is the Cauchy Schwarz inequality, and the third inequality follows from the non-expansiveness of proximal operators:
$$
\left\| X _ {k + 1} - X ^ {*} \right\| = \left\| \operatorname {p r o x} _ {\rho f} \left(Y _ {k}\right) - \operatorname {p r o x} _ {\rho f} \left(Y ^ {*}\right) \right\| \leq \left\| Y _ {k} - Y ^ {*} \right\|.
$$
$$
\begin{array}{l} \operatorname {d i s t} \left\{0, \partial_ {Z} F (X _ {k + 1}, Z _ {k + 1}) \right\} = \operatorname {d i s t} \left\{0, \partial g (X _ {k + 1}) + \frac {1}{\rho} \left(Y ^ {*} - X ^ {*}\right) \right\} \\ \leq \rho^ {- 1} \| \left(X _ {k + 1} - X ^ {*}\right) - \left(Y _ {k} - Y ^ {*}\right) - \left(Y _ {k + 1} - Y _ {k}\right) \| \\ \leq 2 \rho^ {- 1} \| Y _ {k} - Y ^ {*} \| + \rho^ {- 1} \| Y _ {k + 1} - Y _ {k} \| \\ \rightarrow 0. \\ \end{array}
$$
Hence,  $\mathrm{dist}\Big\{0,\partial F(X_{k + 1},Z_{k + 1})\Big\} \to 0$  . This together with (16) and (19), and the fact that  $f$  is partially smooth at  $X^{\star}$  with respect to  $\mathcal{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^\top$ ,  $\phi_0 = \mathbf{0}_m$ , and  $\varphi_0 = \mathbf{0}_n$ , and the default stepsize, the updates in (6) gives that
$$
X _ {k} = \mathbf {0} _ {m \times n}
$$
$$
\phi_ {k} = (k + 1) n ^ {- 1} (p + (m + n) ^ {- 1})
$$
$$
\varphi_ {k} = (k + 1) m ^ {- 1} \left(q + (m + n) ^ {- 1}\right)
$$
when  $k = 1,2,\ldots N$  , where
$$
N = \min _ {i j} \lceil C _ {i j} m n (m + n) ^ {- 1} (m p _ {i} + n q _ {j} + 1) ^ {- 1} - 1 \rceil = O (\max (m, n)).
$$
When the cost is normalized, we can use the rough approximation,  $C_{ij} \sim 1, p_i \sim m^{-1}$ , and  $q_j \sim n^{-1}$  to get the simple initialization strategy:  $X_0 = \mathbf{0}_{m \times n}$ ,  $\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$ , that skips these first  $N$  iterations.
C Derivation of Dual and Stopping criterion
A simple calculation reveals that  $g(X) = \iota_{\mathcal{X}}(X)$  has the Fenchel conjugate
$$
g ^ {*} (U) = \left\{ \begin{array}{l l} \langle U, p q ^ {\top} \rangle , & \text {i f} U = \mu \mathbf {1} _ {n} ^ {\top} + \mathbf {1} _ {m} \nu^ {\top} \\ \infty , & \text {o t h e r w i s e .} \end{array} \right.
$$
Hence, any feasible dual variable is on the form  $U = \mu \mathbf{1}_n^\top + \mathbf{1}_m \nu^\top$ , for which  $g^*(U) = p^\top \mu + q^\top \nu$ . Further, the conjugate of  $f(X) = \langle C, X \rangle + \iota_{\mathbb{R}_+^{m \times n}}(X) + h(X)$  can be expressed as
$$
\begin{array}{l} f ^ {*} (U) = \sup  _ {X} \langle U, X \rangle - f (X) \\ = \sup  _ {X \geq 0} \langle U - C, X \rangle - h (X) \\ \leq \sup  _ {X \geq 0} \left\langle [ U - C ] _ {+}, X \right\rangle - h (X) \\ \leq \sup  _ {X} \left\langle [ U - C ] _ {+}, X \right\rangle - h (X) \\ = h ^ {*} ([ U - C ] _ {+}). \\ \end{array}
$$
Moreover, since  $h$  is sparsity promoting,
$$
\begin{array}{l} h ^ {*} ([ U - C ] _ {+}) = \sup  _ {X} \langle [ U - C ] _ {+}, X \rangle - h (X) \\ \leq \sup  _ {X} \left\langle [ U - C ] _ {+}, [ X ] _ {+} \right\rangle - h ([ X ] _ {+}) \\ = \sup  _ {X \geq 0} \left\langle [ U - C ] _ {+}, X \right\rangle - h (X) \\ = f ^ {*} (U), \\ \end{array}
$$
meaning that  $f^{*}(U) = h^{*}([U - C]_{+})$ . This gives the dual problem
$$
\underset {U = \mu \mathbf {1} _ {n} ^ {\top} + \mathbf {1} _ {m} \nu^ {\top}} {\operatorname * {m a x i m i z e}} - f ^ {*} (- U) - g ^ {*} (U) = - p ^ {\top} \mu - q ^ {\top} \nu - h ^ {*} ([ - \mu \mathbf {1} _ {n} ^ {\top} - \mathbf {1} _ {m} \nu^ {\top} - C ] _ {+}). \tag {20}
$$
To relate the iterates to the optimality condition, we define  $X_{k + 1} = \mathrm{prox}_{\rho f}(Y_k)$ ,  $Z_{k + 1} = \mathrm{prox}_{\rho g}(2X_{k + 1} - Y_k)$ , and  $U_{k} \coloneqq (Y_{k - 1} - X_{k}) / \rho$  and observe that
$$
\rho^ {- 1} (Z _ {k + 1} - X _ {k + 1}) = \rho^ {- 1} (X _ {k + 1} - Y _ {k} - \phi_ {k} {\bf 1} _ {n} - {\bf 1} _ {m} \varphi_ {k} ^ {\top}) = - U _ {k + 1} - (\rho^ {- 1} \phi_ {k}) {\bf 1} _ {n} - {\bf 1} _ {m} (\rho^ {- 1} \varphi_ {k}) ^ {\top}.
$$
By theorem 25.6 in [3],  $Z_{k + 1} - X_{k + 1}\to 0$ , and  $U_{k + 1}\rightarrow U^{\star}$  where  $U^{\star}$  is a solution to (20). Hence  $(\rho^{-1}\phi_k)\mathbf{1}_n + \mathbf{1}_m(\rho^{-1}\varphi_k)^\top \to U^*$ , meaning that  $\rho^{-1}\phi_k$  and  $\rho^{-1}\varphi_k$  will correspond to  $\mu$  and  $\nu$  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:
$$
\begin{array}{l} r _ {\text {p r i m a l}} := \max  (\| X _ {k} \mathbf {1} _ {n} - p \|, \| X _ {k} ^ {\top} \mathbf {1} _ {m} - q \|) = \max  (\| r _ {k} \|, \| s _ {k} \|), \\ \operatorname {g a p} := \left\langle C, X _ {k} \right\rangle + h (X _ {k}) - \rho^ {- 1} \left(p ^ {\top} \phi_ {k} + q ^ {\top} \varphi_ {k}\right) + h ^ {*} \left(\rho^ {- 1} \left(\left[ \phi_ {k} \mathbf {1} _ {n} ^ {\top} + \mathbf {1} _ {m} \varphi_ {k} ^ {\top} - \rho C \right] _ {+}\right) \right. \\ = \langle C, X _ {k} \rangle - \rho^ {- 1} (p ^ {\top} \phi_ {k} + q ^ {\top} \varphi_ {k}) + \Delta_ {h} (X _ {k}, \phi_ {k}, \varphi_ {k}), \\ \end{array}
$$
where  $\Delta_h$  denotes the duality gap deviation compared to the unregularized problem. Further, notice that by letting  $\bar{X}_{k + 1} = [X_k + \phi_k\mathbf{1}_n^\top +\mathbf{1}_m\varphi_k^\top -\rho C]_+$ , we can express  $[\phi_k\mathbf{1}_n^\top +\mathbf{1}_m\varphi_k^\top -\rho C]_+ = [\bar{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.
Table 2: Termination criteria for a selection of regularizers. Here we let  $\parallel  \cdot  \parallel$  denote an arbitrary norm that is sparsity promoting.
h	Δ	r_dual
0	0	||Xk+1-Xk|+|F
||X|_F	[|((1+ρ)Xk+1-Xk|+|F|2	0
||X||	0	||prox_h([\bar{X}_{k+1}-X_k]_+)||_F
D Backpropagation through regularized OT-costs
We let  $h_c = h + \iota_{\mathbb{R}_+^{m \times n}} + \iota_{\mathcal{X}}$ , and define the OT-cost as follows:  $\mathrm{OT}_h(C) = \inf_X \langle C, X \rangle + h_c(X)$ . Note that OT cost is equal to the optimal value of the OT problem (1), i.e. if  $X^{\star}$  is a solution to (1), then  $\mathrm{OT}_h(C) = \langle C, X^{\star} \rangle + h_c(X^{\star})$ . The OT cost is also closely related to the Fenchel conjugate of  $h_c$ , as  $\mathrm{OT}_h(C) = -\bar{h}^*(-C)$ , and hence  $\bar{h}^*(-C) + \bar{h}(X^{\star}) = \langle -C, X^{\star} \rangle$ . Invoking subdifferential properties of the Fenchel conjugate gives that  $X^{\star} \in \partial \bar{h}^*(-C)$ . Further
$$
\partial_ {C} (- \mathrm {O T} _ {h} (C)) = \partial_ {C} h _ {c} ^ {*} (- C) \ni - X ^ {\star}.
$$
In particular, whenever the solution is unique, such as when  $h$  is strongly convex, then  $\nabla_{C}\mathrm{OT}_{h}(C) = X^{\star}$ . 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:
1. The thread is the basic unit of code. All threads execute the same code, but they are endowed with an ID that can be used to parameterize memory access and control decisions.
2. A block is a group of threads that execute together on an SM in the GPU. Threads within a block can cooperate and communicate using shared memory.
3. 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.
1. 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.
2. The shared memory is a low-latency memory space, physically located on the GPU chip. The shared memory is used for data sharing within 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.
3. The register memory of NVIDIA GPUs is the fastest memory with the lowest latency. Registers are private to each thread and are not shared among other threads or blocks. They are used for temporary data storage and computation in a single thread. The access speed for registers is faster than the shared memory, but not by orders of magnitude. Register memory also has limited capacity. NVIDIA Tesla V100 has a limit of 64K 32-bit registers per SM.
E.1.3 Coalesced Memory Access
An important feature of global memory management is coalesced memory access. It is a memory access pattern that maximizes memory bandwidth and improves memory performance in parallel computing, particularly on GPUs. It refers to the way threads in 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 stripped, 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  $\mathrm{Gibth}^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.
Table 3: Comparison of the per-iteration runtime before and after the update for coalesced memory access. Note that the focus here is on the relative change in the per-iteration runtime after decreasing the problem size by one. The improvement in the per-iteration runtime under the same problem size is caused by other optimization techniques.
Problem Size	Runtime Before Update	Runtime After Update
4096 × 4096	0.3577	0.3453
4095 × 4095	0.3812 (+6.58%)	0.3417 (-1.04%)
8192 × 8192	1.281	1.152
8191 × 8191	1.548 (+20.80%)	1.153 (+0.10%)
10240 × 10240	1.916	1.808
10239 × 10239	2.189 (+14.24%)	1.781 (-1.52%)
E.2.2 Quadratically Regularized DROT
Similar to the implementation in [27], the matrices  $X$  and  $C$  are stored in the global memory in column-major order. For the problem sizes that we target, the global memory is the only one that can fit the matrices. In addition, the performance penalty of storing them in a slower memory is limited since each element only needs to be loaded once for every update. As the update of  $X$  is done column by column, the auxiliary variables  $\phi$  and  $\varphi$  can be shared and reused among the threads in a block. To speedup the memory access, we load the required elements of  $\phi$  and  $\varphi$  into shared memory at the beginning of the main kernel.
2The implementation is open-source and available at https://github.com/vienmai/drot
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
(1) Load the corresponding elements of  $\phi$  and  $\varphi$  from the shared memory.
(2) Load the corresponding elements of  $X$  and  $C$  from the global memory. Compute  $[Y_{ij} - \rho C_{ij}]_+$  and  $C_{ij} \cdot X_{ij}$ . For the quadratically regularized case, we can directly update  $X_{ij}$  with  $Y_{ij}$  after multiplying it with the scale  $1 / (1 + \rho \alpha)$ .
(3) Perform in-warp reduction to compute the column sum and the objective value in the block.
(4) Use atomic operations to accumulate the results to the corresponding row and column sums.
(5) This step is only used by the group-lasso kernel, introduced in Section E.2.3.
Similar to the implementation by [27], the iterations are divided into even iterations and odd iterations.
1. For the even iterations, the kernel updates  $X_{k + 1}$  as
$$
X _ {k + 1} = \operatorname {p r o x} _ {\rho h} \left(\left[ X _ {k} + \phi_ {k + 1} \mathbf {1} _ {n} ^ {\top} + \mathbf {1} _ {m} \varphi_ {k + 1} ^ {\top} - \rho C \right] _ {+}\right) - \rho C
$$
2. For the odd iterations, it updates  $X_{k + 1}$  as
$$
X _ {k + 1} = \operatorname {p r o x} _ {\rho h} \left(\left[ X _ {k + 1} + \phi_ {k + 1} \mathbf {1} _ {n} ^ {\top} + \mathbf {1} _ {m} \varphi_ {k + 1} ^ {\top} \right] _ {+}\right)
$$
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:
(1) Load the corresponding elements of  $\phi$  and  $\varphi$  from the shared memory.
(2) Load the corresponding elements of  $X$  and  $C$  from the global memory. Compute  $[Y_{ij} - \rho C_{ij}]_+$  and  $C_{ij} \cdot 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).
(3) Perform in-warp reduction. For the group-lasso regularized case, we also reduce the sum of  $[Y_{ij} - \rho C_{ij}]_+^2$  to compute the norm of the group.
(4) Accumulate the results to the corresponding sum of the row, the sum of the column, and the sum of the squared elements within the group using atomic operations.
(5) Compute the scale for the group and apply it to the computed  $[Y_{ij} - \rho 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
Figure 5: GPU Kernel Design