Markovian Sliced Wasserstein Distances: Beyond Independent Projections

We first start with a general definition of Markovian sliced Wasserstein distance in Definition 1.

The first projecting direction $\theta_{1}$ follows the distribution $\sigma_{1}(\theta_{1})$ with $\sigma_{1}(\theta_{1})$ to be any distributions on the unit hyper-sphere, e.g., the uniform distribution, a von Mises-Fisher distribution, and so on. By designing the transition distribution $\sigma_{l}(\theta_{l}|\theta_{l-1})$, we can obtain various variants of MSW. It is worth noting that the MSW can be rewritten as the average of $T$ expectation of one-dimensional Wasserstein distance, $\text{MSW}^{p}_{p,T}(\mu,\nu)=\frac{1}{T}\sum_{t=1}^{T}\mathbb{E}_{\theta_{t}\sim\sigma_{t}(\theta_{t})}[W_{p}^{p}\left(\theta_{t}\sharp\mu,\theta_{t}\sharp\nu\right)]$, however, $\sigma_{t}(\theta_{t})=\int\prod_{i=1}^{t}\sigma_{1}(\theta_{1})\prod_{l=2}^{t}\sigma_{t^{\prime}}(\theta_{t^{\prime}}|\theta_{t^{\prime}-1})d\theta_{1}\ldots d\theta_{t-1}$ are not the same for $t=1,\ldots,T$. Moreover, sampling directly from $\sigma_{t}(\theta_{t})$ is intractable, hence, using the definition of the MSW in Definition 1 is more reasonable in terms of approximating the expectation using Monte Carlo samples.

Monte Carlo estimation. Similar to SW, we also need to use Monte Carlo samples to approximate the expectation in Definition 1. We first samples $\theta_{11},\ldots,\theta_{L1}\sim\sigma_{1}(\theta_{1})$ for $L\geq 1$, then we samples $\theta_{lt}\sim\sigma_{t}(\theta_{t}|\theta_{lt-1})$ for $t=1,\ldots,T$ and $l=1,\ldots,L$. After that, we can form an unbiased empirical estimation of MSW as follows: $\widehat{\text{MSW}}^{p}_{p,T}(\mu,\nu)=\frac{1}{LT}\sum_{l=1}^{L}\sum_{t=1}^{T}\text{W}_{p}^{p}(\theta_{lt}\sharp\mu,\theta_{lt}\sharp\nu).$

Before going to the specific design of those distributions, we first discuss the empirical estimation of MSW, and investigate its theoretical properties including topological properties, statistical properties, and computational properties.

Topological Properties. We first state the following assumption: A1. In MSW, the prior distribution $\sigma_{1}(\theta_{1})$ is supported on all the unit-hypersphere or there exists a transition distribution $\sigma_{t}(\theta_{t}|\theta_{t-1})$ being supported on all the unit-hypersphere. The assumption A1 is easy to satisfy and it holds for all later choices of the prior distribution and transition distribution. We now consider the metricity properties of the Markovian sliced Wasserstein distance.

The proof of Theorem 1 is in Appendix B.1. Next, we show that the convergence in MSW implies the weak convergence of probability measures and the reverse also holds.

Theorem 2 means that for any sequence of probability measures $(\mu_{k})_{k\in\mathbb{N}}$ and $\mu$ in $\mathcal{P}_{p}(\mathbb{R}^{d})$, $\lim_{k\to+\infty}\text{MSW}_{p,T}(\mu_{k},\mu)=0$ if and only if for any continuous and bounded function $f:\mathbb{R}^{d}\to\mathbb{R}$, $\lim_{k\rightarrow+\infty}\int f\mathrm{~{}d}\mu_{k}=\int f\mathrm{~{}d}\mu$. The proof of Theorem 2 is in Appendix B.2. Next, we discuss the connection of MSW to previous sliced Wasserstein variants.

The proof of Proposition 1 is in Appendix B.3.

Statistical Properties. We investigate the sample complexity (empirical estimation rate) of the MSW.

The proof of Proposition 2 is in Appendix B.4. The proposition suggests that MSW does not suffer from the curse of dimensionality. Next, we investigate the MSW’s Monte Carlo approximation error.

The proof of Proposition 3 is in Appendix B.5. From the above proposition, we know that increasing the number of projections $L$ reduces the approximation error.

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_{2}n+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.

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 $\mathcal{S}^{d-1}_{\theta_{t-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}(\mathcal{S}^{d-1}_{\theta_{t-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: $\text{Proj}_{\theta_{t-1}}(\theta_{t})=\text{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 $\text{Proj}_{\mathbb{S}^{d-1}}(\theta)=\frac{\theta}{||\theta||_{2}}$ is the normalizing operator. In a greater detail, we first sample $\theta^{\prime}_{t}\sim\mathcal{U}(\mathbb{S}^{d-1})$ and then set $\theta_{t}=\text{Proj}_{\theta_{t-1}}(\theta^{\prime}_{t})$. Therefore, we have $\sigma_{t}(\theta_{t}|\theta_{t-1})=\mathcal{U}(\mathcal{S}^{d-1}_{t-1})=\text{Proj}_{\theta_{t-1}}\sharp\mathcal{U}(\mathbb{S}^{d-1})$.

Input-awared 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)=\text{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})=\text{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}\text{vMF}(\theta|\epsilon,\kappa)=\mathcal{U}(\mathbb{S}^{d-1})$ and $\lim_{\kappa\to\infty}\text{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\ldots\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-awared 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 maximas of the optimization problem $\max_{\theta\in\mathbb{S}^{d-1}}\text{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-awared vMF transition, it is also unclear if the stationary distribution exists when the parameter $\kappa<\infty$.

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.

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_{2}n+(T-M)Ldn)/N)$ in the orthogonal-based transitions. In the case of the input-awared transition, the computational complexity is still $\mathcal{O}(TLn\log_{2}n+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-awared 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.

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)=\left(X_{i}(t)\right)_{i=1}^{n}\in\left(\mathbb{R}^{d}\right)^{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 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,\ldots,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.

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\}$)(Max-K-SW (K=1) is Max-SW), MSW (all variant, $L=\{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.