SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Computational Properties. When  $\mu$  and  $\nu$  are two discrete probability measures in  $\mathcal{P}_p(\mathbb{R}^d)$  that have at most  $n$  supports, the computational complexity for the Monte Carlo approximation of MSW is  $\mathcal{O}(TLn\log_2n + TLdn)$  where  $\mathcal{O}(TLn\log n)$  is for computation of  $TL$  one-dimensional Wasserstein distances and  $\mathcal{O}(TLdn)$  is the projecting complexity for  $TL$  projections from  $d$  dimension to 1 dimension. The memory complexity of MSW is  $\mathcal{O}(TL(d + n))$  for storing the projecting directions and the projections.
3.2 Specific Choices of Projecting Distributions
Designing the projecting distribution  $\sigma(\theta_1, \ldots, \theta_T)$  is the central task in using MSW since it controls the projecting behavior. For each choice of the  $\sigma(\theta_1, \ldots, \theta_T)$ , we obtain a variant of MSW. Since we impose the first order Markov structure  $\sigma(\theta_1, \ldots, \theta_T) = \sigma_1(\theta_1) \prod_{t=2}^{T} \sigma_t(\theta_t | \theta_{t-1})$ , there are two types of distributions that we need to choose: the prior distribution  $\sigma_1(\theta_1)$  and the transition distribution  $\sigma_t(\theta_t | \theta_{t-1})$  for all  $t = 2, \ldots, T$ .
Prior distribution. The most simple choice of  $\sigma_{1}(\theta_{1})$  when we know nothing about probability measures that we want to compare is the uniform distribution over the unit hypersphere  $\mathcal{U}(\mathbb{S}^{d - 1})$ . Moreover, the metricity of MSW is guaranteed regardless of the transition distribution with this choice. Therefore, the uniform distribution is the choice that we use in our experiments in the paper. It is worth noting that we could also use a distribution that is estimated from two interested probability measures [44]; however, this approach costs more computation.
Transition distribution. We discuss some specific choices of the transition distributions  $\sigma_t(\theta_t|\theta_{t - 1})$ . Detailed algorithms for computing MSW with specific transitions are given in Appendix A.3.
Orthogonal-based transition. Motivated by the orthogonality constraint in Max-K-SW and K-SW, we can design a transition distribution that gives us an orthogonal projecting direction to the previous one. In particular, given a previous projecting direction  $\theta_{t-1}$ , we want to have  $\theta_t$  such that  $\langle \theta_t, \theta_{t-1} \rangle = 0$ , namely, we want to sample from the subsphere  $S_{\theta_{t-1}}^{d-1} := \{\theta_t \in \mathbb{S}^{d-1} | \langle \theta_t, \theta_{t-1} \rangle = 0\}$ . To the best of our knowledge, there is no explicit form of distribution (known pdf) that is defined on that set. However, we can still sample from the uniform distribution over that set:  $\mathcal{U}(S_{\theta_{t-1}}^{d-1})$  since that distribution can be constructed by pushing the uniform distribution over the whole unit hypersphere  $\mathcal{U}(\mathbb{S}^{d-1})$  through the projection operator:  $\mathrm{Proj}_{\theta_{t-1}}(\theta_t) = \mathrm{Proj}_{\mathbb{S}^{d-1}}\left(\theta_t - \frac{\langle \theta_{t-1}, \theta_t \rangle}{\langle \theta_{t-1}, \theta_{t-1} \rangle} \theta_{t-1}\right)$  where  $\mathrm{Proj}_{\mathbb{S}^{d-1}}(\theta) = \frac{\theta}{||\theta||_2}$  is the normalizing operator. In a greater detail, we first sample  $\theta_t' \sim \mathcal{U}(\mathbb{S}^{d-1})$  and then set  $\theta_t = \mathrm{Proj}_{\theta_{t-1}}(\theta_t')$ . Therefore, we have  $\sigma_t(\theta_t | \theta_{t-1}) = \mathcal{U}(S_{t-1}^{d-1}) = \mathrm{Proj}_{\theta_{t-1}}\# \mathcal{U}(\mathbb{S}^{d-1})$ .
Input-aware transition. The above two transition distributions do not take into account the information of the two probability measures  $\mu$  and  $\nu$  that we want to compare. Hence, it could be inefficient to explore good projecting directions by comparing  $\mu$  and  $\nu$ . Motivated by the projected sub-gradient ascent [9] update in finding the "max" projecting direction, we could design the transition distribution as follows:  $\sigma_t(\theta_t|\theta_{t-1}) = \delta_{f(\theta_{t-1}|\eta, \mu, \nu)}$  where  $\delta$  denotes the Dirac Delta function and the transition function  $f(\theta_{t-1}|\eta, \mu, \nu) = \mathrm{Proj}_{\mathbb{S}^{d-1}}\left(\theta_{t-1} + \eta \nabla_{\theta_{t-1}} W_p\left(\theta_{t-1}\sharp \mu, \theta_{t-1}\sharp \nu\right)\right)$ , with  $\eta > 0$  is the stepsize hyperparameter. As the current choice is a deterministic transition, it requires the prior
distribution to have supports on all  $\mathbb{S}^{d - 1}$  to obtain the metricity for MSW. A choice to guarantee the metricity regardless of the prior distribution is the von Mises-Fisher (vMF) distribution [23], namely,  $\sigma_t(\theta_t|\theta_{t - 1}) = \mathrm{vMF}(\theta_t|\epsilon = f(\theta_{t - 1}|\eta ,\mu ,\mu),\kappa)$ . The details about the vMF distribution including its probability density function, its sampling procedure, and its properties are given in Appendix A.2. In summary, the vMF distribution has two parameters: the location parameter  $\epsilon \in \mathbb{S}^{d - 1}$  which is the mean, and the concentration parameter  $\kappa \in \mathbb{R}^+$  which plays the role as the precision. Thank the interpolation properties of the vMF distribution:  $\lim_{\kappa \to 0}\mathrm{vMF}(\theta |\epsilon ,\kappa) = \mathcal{U}(\mathbb{S}^{d - 1})$  and  $\lim_{\kappa \to \infty}\mathrm{vMF}(\theta |\epsilon ,\kappa) = \delta_{\epsilon}$ , the transition distribution can balance between heading to the "max" projecting direction and exploring the space of directions.
Stationarity of  $\sigma_T(\theta_T)$ . A natural important question arises: what is the distribution of  $\sigma_T(\theta_T) = \int \dots \int \sigma(\theta_1, \ldots, \theta_T)d\theta_1 \ldots d\theta_{T-1}$  when  $T \to \infty$ ? The answer to the above questions depends on the choice of the projection distribution which is discussed in Section 3.2. For the Orthogonal-based transitions and the uniform distribution prior, it is unclear whether the stationary distribution exists. For the deterministic Input-aware transition and the uniform prior, we have  $\lim_{T \to \infty} \sigma_T(\theta_T) = \sum_{a=1}^{A} \alpha_a \delta_{\theta_a^*}$  with  $\sum_{a=1}^{A} \alpha_a = 1$  where  $\theta_a^* (a = 1, \ldots, A)$  are local maxima of the optimization problem  $\max_{\theta \in \mathbb{S}^{d-1}} \mathrm{W}_p(\theta \sharp \mu, \theta \sharp \nu)$  and some unknown weights  $\alpha_a$  that depend on  $\mu$  and  $\nu$ . This property is due to the fact that the projected sub-gradient ascent can guarantee local maxima convergence [46]. For the Input-aware vMF transition, it is also unclear if the stationary distribution exists when the parameter  $\kappa < \infty$ .
3.3 Burning and Thinning
In the definition of MSW in Definition 1, we take the expectation on the joint distribution over all timesteps  $\sigma(\theta_{1:T})$  which leads to the time and memory complexities to be linear with  $T$  in the Monte Carlo approximation. Therefore, we can adapt the practical technique from MCMC methods which is burning and thinning to reduce the number of random variables while still having the dependency.
Definition 2. For any  $p \geq 1$ ,  $T \geq 1$ , dimension  $d \geq 1$ , the number of burned steps  $M \geq 0$ , and the number of thinned steps  $N \geq 1$ , the burned thinned Markovian sliced Wasserstein of order  $p$  between two probability measures  $\mu \in \mathcal{P}_p(\mathbb{R}^d)$  and  $\nu \in \mathcal{P}_p(\mathbb{R}^d)$  is:
$$
M S W _ {p, T, N, M} ^ {p} (\mu , \nu) = \mathbb {E} \left[ \frac {N}{T - M} \sum_ {t = 1} ^ {(T - M) / N} W _ {p} ^ {p} \left(\theta_ {t} ^ {\prime} \sharp \mu , \theta_ {t} ^ {\prime} \sharp \nu\right) \right], \tag {8}
$$
where the expectation is under the projection distribution  $\theta_{1:N(T - M)}^{\prime} \sim \sigma(\theta_{1:N(T - M)}^{\prime})$  with  $\sigma(\theta_{1:N/(T - M)}^{\prime})$  being the marginal distribution which is obtained by integrating out random projecting directions at the time step  $t$  such that  $t \leq M$  or  $t\% N \neq 0$  from  $\sigma(\theta_1, \ldots, \theta_T) = \sigma_1(\theta_1) \prod_{l=2}^{T} \sigma_t(\theta_t | \theta_{t-1})$  for all  $t = 1, \ldots, T$ .
Similar to MSW, the burned-thinned MSW is also a metric on  $\mathcal{P}_p(\mathbb{R}^d)$  when there exists a time step  $t$  that is not burned, is not thinned, and  $\theta_t$  is a random variable that has the supports on all  $\mathbb{S}^{d-1}$ . We discuss more details about the burned-thinned MSW including its topological and statistical properties in Appendix A.4. The Monte Carlo estimation of the burned-thinned MSW is given in Equation 10 in Appendix A.4. The approximation is the average of the projected Wasserstein distance from  $\theta_{tl}$  with  $t \geq M$  and  $t\% N = 0$ . By reducing the number of random projecting directions, the computational complexity of the burned-thinned MSW is improved to  $\mathcal{O}((T - M)Ln\log_2n + (T - M)Ldn) / N)$  in the orthogonal-based transitions. In the case of the input-aware transition, the computational complexity is still  $\mathcal{O}(TLn\log_2n + TLdn)$  since the transition requires computing the gradient of the projected Wasserstein distance. However, in all cases, the memory complexity is reduced to  $\mathcal{O}((T - M)L(d + n) / N)$ .
Burned thinned MSW is the generalization of Max-SW. The empirical computation of Max-SW [14] with the projected sub-gradient ascent and uniform random initialization can be viewed as a special case of burned thinned MSW with the input-aware transition and with the number of burned samples  $M = T - 1$ . The difference is that Max-SW uses only one local maximum to compute the distance while the burned thinned MSW uses  $L \geq 1$  maximums (might not be unique).
More discussions. We refer the reader to Appendix A.5 for other related discussions e.g., "K-SW is autoregressive decomposition of projecting distribution", "sequential generalization of Max-K-SW", and related literature.
Figure 1: The figures show the gradient flows that are from the empirical distribution over the color points to the empirical distribution over S-shape points. The corresponding Wasserstein-2 distance between the empirical distribution at the current step and the S-shape distribution and the computational time (in seconds) to reach the step is reported at the top of the figure.
Figure 2: The figures show the source image, the target image, and the transferred images from different distances. The corresponding Wasserstein-2 distance between the empirical distribution over transferred color palates and the empirical distribution over the target color palette and the computational time (in second) are reported at the top of the figure.
In this section, we refer MSW with orthogonal-based transition as oMSW, MSW with input-aware transition as iMSW (using the Dirac distribution) and viMSW (using the vMF distribution). We compare MSW variants to SW, Max-SW, K-SW, and Max-K-SW in standard applications e.g., gradient flows, color transfer, and deep generative models. Moreover, we also investigate the role of hyperparameters, e.g., concentration parameter  $\kappa$ , the number of projections  $L$ , the number of time steps  $T$ , the number of burning steps  $M$ , and the number of thinning steps  $N$  in applications.
4.1 Gradient Flows and Color Transfer
Gradient flows. We follow the same setting in [17]. The gradient flow models a distribution  $\mu(t)$  flowing with time  $t$  along the gradient flow of a loss functional  $\mu(t) \to \mathcal{D}(\mu(t), \nu)$  that drives it towards a target distribution  $\nu$  [53] where  $\mathcal{D}$  is a given distance between probability measures. In this setup, we consider  $\nu = \frac{1}{n} \sum_{i=1}^{n} \delta_{Y_i}$  as a fixed empirical target distribution and the model distribution  $\mu(t) = \frac{1}{n} \sum_{i=1}^{n} \delta_{X_i(t)}$ . Here, the model distribution is parameterized by a time-varying point cloud  $X(t) = (X_i(t))_{i=1}^n \in (\mathbb{R}^d)^n$ . Starting from an initial condition at time  $t = 0$ , we integrate the ordinary differential equation  $\dot{X}(t) = -n \nabla_{X(t)} \left[ \mathcal{D} \left( \frac{1}{n} \sum_{i=1}^{n} \delta_{X_i(t)}, \nu \right) \right]$  for each iteration. In the experiments, we utilize the Euler scheme with 300 timesteps and the step size is  $10^{-3}$  to move the empirical distribution over colorful 1000 points  $\mu(0)$  to the distribution over S-shape 1000 points  $(\nu)$  (see Figure 1). For Max-SW, Max-K-SW, iMSW, and vIMSW, we use the learning rate parameter for projecting directions  $\eta = 0.1$ . We report the Wasserstein-2 distances between the empirical distribution  $\mu(t)$  and the target empirical distribution  $\nu$ , and the computational time in Table 1. We also give the visualization of some obtained flows in Figure 1. We refer the reader to Figure 4 in Appendix C.1 for the full visualization of all flows and detailed algorithms. We observe that iMSW
Table 1: Summary of Wasserstein-2 distances, computational time in second (s) of different gradient flows.
Distances	Wasserstein-2 (↓)	Time (↓)	Distances	Wasserstein-2 (↓)	Time (↓)
SW (L=30)	0.0099 × 10-2	1.55	Max-K-SW (K=2,T=15)	0.0146 × 10-2	3.35
Max-SW (T=30)	0.0098 × 10-2	3.48	iMSW (L=2,T=5) (ours)	0.0064 × 10-2	1.41
K-SW (L=15,K=2)	0.0098 × 10-2	1.71	viMSW (L=2,T=5,κ=50)(ours)	0.0043 × 10-2	2.94
Table 2: Summary of FID and IS scores from three different runs on CIFAR10 (32x32), and CelebA (64x64).
Method	CIFAR10 (32x32)	CelebA (64x64)	Method	CIFAR10 (32x32)	CelebA (64x64)
FID (↓)	IS (↑)	FID (↓)	FID (↓)	IS (↑)	FID (↓)
SW	14.21±1.12	8.19±0.07	8.93±0.23	oMSW (ours)	14.12±0.54	8.20±0.05	9.68±0.55
Max-K-SW	14.38±0.08	8.15±0.02	8.94±0.35	iMSW (ours)	14.12±0.48	8.24±0.09	8.89±0.23
K-SW	15.24±0.02	8.15±0.03	9.41±0.16	viMSW (ours)	13.98±0.59	8.12±0.20	8.91±0.11
gives better flows than SW, Max-SW, K-SW, and Max-K-SW. Namely, the empirical distribution  $\mu(t) (t = 300)$  with iMSW is closer to  $\nu$  in terms of Wasserstein distance. More importantly, iMSW consumes less computation than its competitors since it can use a smaller number of projections due to more informative projecting directions. Furthermore, viMSW gives better final results than iMSW, however, the trade-off is doubling the time computation due to the sampling step of vMF distribution. In this case, K-SW is equivalent to our oMSW with  $T = K = 2$  since the dimension  $d = 2$ . We refer the reader to Appendix C.1 for more discussion.
Studies on hyperparameters. From Table 3 in Appendix C.1, increasing the number of projections  $L$  yields better performance for SW, K-SW, and iMSW. Similarly, increasing the number of timesteps  $T$  also helps Max-SW and iMSW better. Moreover, we find that for the same number of total projections e.g.,  $L = 5, T = 2$  and  $T = 2, L = 5$ , a larger timestep  $T$  might lead to a better result for iMSW. For burning and thinning, we see that they help to reduce the computation while the performance stays comparable or even better if choosing the right value of  $M$  and  $N$ . Also, iMSW the burning steps  $M = T - 1$  is still better than Max-SW with  $T$  time steps. For the concentration parameter  $\kappa$  in viMSW, the performance of viMSW is not monotonic in terms of  $\kappa$ .
Color transfer. We aim to transfer the color palate (RGB) of a source image to the color palette (RGB) target image. Therefore, it is natural to build a gradient flow that starts from the empirical distribution over the color palette of the source image to the empirical distribution over the color palette of the target image. Since the value of color palette is in the set  $\{0, \dots, 255\}^3$ , we round the value of the supports of the empirical distribution at the final step of the Euler scheme with 2000 steps and  $10^{-3}$  step size. Greater detail can be found in Appendix C.2. For Max-SW, Max-K-SW, iMSW, and viMSW, we use the learning rate parameter for projecting directions  $\eta = 0.1$ . We show the transferred images, the corresponding Wasserstein-2 distances between the empirical distribution over the transferred color palette and the empirical distribution over the target color palette, and the corresponding computational time in Figure 2. From the figures, iMSW and viMSW give the best transferred images quantitatively and qualitatively. Moreover, oMSW is comparable to SW, Max-SW, K-SW, and are better than Max-K-SW while consuming much less computation. We refer the reader to Figure 5 in Appendix C.2 for the color palette visualization and to Figure 6 for another choice of the source and target images. We also conduct studies on hyperparameters in Appendix C.2 where we observe some similar phenomenons as in gradient flow.
4.2 Deep Generative Models
We follow the setup of sliced Wasserstein deep generative models in [15]. The full settings of the framework including neural network architectures, training framework, and hyperparameters are given Appendix C.3. We compare MSW with previous baselines including SW, Max-SW, K-SW, and Max-K-SW on benchmark datasets: CIFAR10 (image size 32x32) [29], and CelebA [36] (image size 64x64). The evaluation metrics are FID score [21] and Inception score (IS) [51] (except on CelebA since IS score poorly captures the perceptual quality of face images [21]). A notable change in computing Max-SW is that we do not use momentum in optimization for max projecting direction like in previous works [26, 42], which leads to a better result.
Summary of generative performance. We train generative models with SW  $(L \in \{100, 1000, 10000\})$ , K-SW  $(L \in \{1, 10, 100\}, K = 10)$ , Max-K-SW  $(K = \{1, 10\}, T = \{1, 10, 100, 1000\}, \eta \in \{0.01, 0.1\})(\text{Max-K-SW (K=1) is Max-SW}), \text{MSW (all variant, } L =$
Figure 3: The FID scores and the IS scores over epochs, and some generated images from CelebA.
$\{10, 100\}, T \in \{10, 100\}$ , iMSW and viMSW ( $\eta \in \{0.01, 0.1\}$ ), and viMSW and ( $\kappa \in \{10, 50\}$ ). We report the best FID score and the best IS score for each distance in Table 2. In addition, we show how scores change with respect to the training epochs in Figure 3. Overall, we observe that viMSW and iMSW give the best generative performance in terms of the final scores and fast convergence on CIFAR10 and CelebA. The oMSW gives comparable results to baselines. Since most computation in training deep generative models is for updating neural networks, the computational time for distances is almost the same. Furthermore, we show some generated images on CelebA in Figure 3 and all generated images on CIFAR10 and CelebA in Figure 7 and Figure 8 in Appendix C.3. We visually observe that the qualitative results are consistent with the quantitative results in Table 2.
Studies on hyperparameters. We conduct experiments to understand the behavior of the burning and thinning technique, and to compare the role of  $L$  and  $T$  in Table 5 in Appendix C.3. Overall, burning (thinning) sometimes helps to improve the performance of training generative models. There is no clear sign of superiority between burning and thinning. We compare two settings of the same number of total projections (same complexities):  $L = 10, T = 100$  and  $L = 100, T = 10$ . On CIFAR10, the first setting is better while the reverse case happens on CelebA.
We have introduced the Markovian sliced Wasserstein (MSW), a novel family of sliced Wasserstein (SW) distances, which imposes a first-order Markov structure on projecting directions. We have investigated the theoretical properties of MSW including topological properties, statistical properties, and computational properties. Moreover, we have discussed three types of transition distribution for MSW, namely, orthogonal-based, and input-aware transitions. In addition, we have proposed a burning and thinning technique to improve the computational time and memory of MSW. Finally, we have compared MSW to previous variants of SW in gradient flows, color transfer, and generative modeling to show that MSW distances are both effective and efficient.
[59] M. J. Wainwright. High-dimensional statistics: A non-asymptotic viewpoint. Cambridge University Press, 2019.
[60] J. Wu, Z. Huang, D. Acharya, W. Li, J. Thoma, D. P. Paudel, and L. V. Gool. Sliced Wasserstein generative models. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pages 3713-3722, 2019.
[61] M. Yi and S. Liu. Sliced Wasserstein variational inference. In Fourth Symposium on Advances in Approximate Bayesian Inference, 2021.
Supplement to "Markovian Sliced Wasserstein Distances: Beyond Independent Projections"
In this supplementary material, we present additional materials in Appendix A. In particular, we provide additional background on sliced Wasserstein variants in Appendix A.1, background on von Mises-Fisher distribution in Appendix A.2, algorithms for computing Markovian sliced Wasserstein distances in Appendix A.3, additional information about burned thinned MSW in Appendix A.4, and discussion on related works in Appendix A.5. We then provide skipped proofs in the main paper in Appendix B. Additional experiments are presented in Appendix C.
A Additional Materials
A.1 Background on Sliced Wasserstein Variants
We review computational aspects of sliced Wasserstein variants.
Computation of Max sliced Wasserstein distance. We demonstrate the empirical estimation of Max-SW via projected sub-gradient ascent algorithm in Algorithm 1. The initialization step for  $\hat{\theta}_0$  is rarely discussed in previous works. Normally,  $\hat{\theta}_0$  is randomly initialized by drawing from the uniform distribution over the unit-hypersphere. Many previous works [26, 44, 45, 42] use Adam update instead of the standard gradient ascent update for Max-SW. In this work, we find out that using the standard gradient ascent update is more stable and effective.
Input. Probability measures  $\mu, \nu$ , learning rate  $\eta$ , the order  $p$ , and the number of iterations  $T$ .
Initialize  $\hat{\theta}_0$
for  $t = 1$  to  $T - 1$  do
$$
\hat {\theta} _ {t} = \hat {\theta} _ {t - 1} + \eta \cdot \nabla_ {\hat {\theta} _ {t - 1}} \mathbf {W} _ {p} (\hat {\theta} _ {t - 1} \sharp \mu , \hat {\theta} _ {t - 1} \sharp \nu)
$$
$$
\hat {\theta} _ {t} = \frac {\hat {\theta} _ {t}}{| | \hat {\theta} _ {t} | | _ {2}}
$$
Return.  $\mathbf{W}_p(\hat{\theta}_T\sharp \mu ,\hat{\theta}_T\sharp \nu)$
K sliced Wasserstein distance. We first review the Gram-Schmidt process in Algorithm 2. With the Gram-Schmidt process, the sampling from  $\mathcal{U}(\mathbb{V}_K(\mathbb{R}^d))$  can be done by sampling  $\theta_1,\ldots ,\theta_k$  i.i.d from  $\mathcal{N}(0,I_d)$  then applying the Gram-Schmidt process on them. Therefore, we present the computation of K sliced Wasserstein distance in Algorithm 3. We would like to recall that the original work of K-SW [50] uses only one set of orthogonal projecting directions. Here, we generalize the original work by using  $L$  sets of orthogonal projecting directions.
Algorithm 1 Max sliced Wasserstein distance
Algorithm 2 Gram-Schmidt process
Input.  $K$  vectors  $\theta_{1},\ldots ,\theta_{K}$ $\theta_{1} = \frac{\theta_{1}}{||\theta_{1}||_{2}}$    
for  $k = 2$  to  $K$  do for  $i = 1$  to  $k - 1$  do  $\theta_{k} = \theta_{k} - \frac{\langle\theta_{i},\theta_{k}\rangle}{\langle\theta_{i},\theta_{i}\rangle}\theta_{i}$  end for  $\theta_{k} = \frac{\theta_{k}}{||\theta_{k}||_{2}}$    
end for   
Return.  $\theta_{1},\dots ,\theta_{K}$
Max K sliced Wasserstein distance. We now present the empirical estimation of Max-K-SW via projected sub-gradient ascent algorithm in Algorithm 4. This algorithm is first discussed in the original paper of Max-K-SW [12]. The optimization of Max-K-SW can be solved by using Riemannian optimization since the Stiefel manifold is a Riemannian manifold. However, to the best of our knowledge, Riemannian optimization has not been applied to Max-K-SW.
Algorithm 3 K sliced Wasserstein distance
Input. Probability measures  $\mu, \nu$ , the dimension  $d$ , the order  $p$ , the number of projections  $L$ , the number of orthogonal projections  $K$ .
Draw  $\theta_{l1},\ldots ,\theta_{lK}$  i.i.d from  $\mathcal{N}(0,I_d)$