Sliced Wasserstein Estimation with Control Variates

We would like to thank the reviewer for the insightful comments. We would like to discuss them as follow:

Q14: The paper lacks positioning with respect to competitors whose goal is also to improve the sampling scheme of SW, or more generally, other SW variants. Evaluating the proposed methods among numerous competitors can be challenging. However, CV-SW is benchmarked only against SW, whereas algorithms such as max-SW (which considers only the 'max' direction), distributional SW (which searches for an 'optimal' direction), or even projected Wasserstein (which uses orthogonal directions) are closely related to the proposed method in spirit. Theoretical and experimental comparisons are missing.

A14: Thank you very much for your insightful comments. Max sliced Wasserstein (Max-SW), distributional sliced Wasserstein (DSW), and Projected Wasserstein (PW) are variant of Sliced Wasserstein which focuses on changing the slicing distribution of SW from uniform distribution over the sphere to a different distribution (e.g., optimal Dirac delta measure, optimal von-Mises Fisher distribution, and so on). They are different distances that are different from SW i.e. different topological properties, different computational schemes. In contrast, we focus on the numerical approximation/estimation aspect of SW by reducing the variance of the estimator. Our control variates are constructed specifically for the uniform slicing distribution.

Although our control variate are constructed for the uniform distribution, they are still applicable for DSW, and PW by using importance sampling. In more detail, we can rewrite $DSW(\theta) = = \mathbb{E}_{U}[W(\theta)\frac{\sigma(\theta)}{\sigma_0 (\theta)}],$ For $\sigma(\theta)$ is the slicing distribution from DSW and $\sigma_0(\theta) = \mathcal{U}(\mathbb{S}^{d-1})$ is the uniform distribution. Therefore, we can use the proposed control variates to reduce the variance of the importance sampling. For PW, we can write it as:

$PW(\theta) = = \mathbb{E}_{\sigma_0(\theta_1)\ldots \sigma_0(\theta_K)}[W(\theta_1,\ldots,\theta_K)\frac{\sigma(\theta_1,\ldots,\theta_K))}{\sigma_0(\theta_1)\ldots \sigma_0(\theta_K)}],$ For $\sigma(\theta_1,\ldots,\theta_K)$ is the orthogonal-based slicing distribution from PW. Therefore, we can still apply the control variates for PW.

Finally, we would like to recall distances such as DSW, PW, and Max-SW are more computationally expensive than SW. Therefore, SW has been still widely used in applications.

Q15: how does the method compares with max-sliced SW, distributional SW and projected SW?

A15: As discussed in Q14, SW is easy and stable to compute since it does not involve optimization as DSW, Max-SW, and orthogonal constraints as PW. From the statistical perspective, the uniform slicing distribution is suitable to use when we do not have any prior information about two compared distributions. Moreover, LCV-SW and UCV-SW are estimations of SW which means that they are still SW. The main aim of the paper is to bridge the literature on variance reduction and SW by proposing control variates for SW. These control variates can still be applied to DSW and PW as discussed.

Q16:Is there any results that compare the value provided by CV-SW with SW or Wasserstein ?

A16: We would like to recall that we do not propose a new SW metric. We propose new estimators for SW. Therefore, the interested distance is still SW. CV-SW will asymptotically converge to SW when increasing the number of Monte Carlo samples or the number of projections. Since the projecting direction belongs to the unit-hypersphere, we can show SW is a lower bound of Wasserstien distance. Therefore, CV-SW might be a lower-bound of Wasserstein distance with high probability.

Q17: Do you have any results regarding the sample complexity?

A17: We are still interested in SW which does not suffer from the curse of dimensionality and has a well-investigated sample complexity [R1] [R2] i.e., the expected value of SW between a measure and its corresponding empirical measure. Since CV-SW is an estimation of SW, the sample complexity has not been defined for CV-SW. For estimation, the interested quantity is an approximation error. As mentioned by Reviewer pTc9, the theoretical understanding of approximation error of using control variate is limited. However, it is widely recognized that using control variate is a practical method for improving the approximation error by reducing the variance.

[R1] Statistical and Topological Properties of Sliced Probability Divergences

[R2] Statistical, Robustness, and Computational Guarantees for Sliced Wasserstein Distances

We would like to thank the reviewer again for giving us additional feedback.

We would like to clarify our claim of "PW, and Max-SW are more computationally expensive than SW".

For Max-SW, a gradient ascent algorithm needs to be used to find the max projecting direction. To the best of our knowledge, there is no guarantee that the algorithm can converge to the global maxima. In practice, the algorithm is used with $T$ iterations. For each iteration, the gradient is computed based on the value of the one-dimensional Wasserstein distance, therefore the overall computational complexity is $\mathcal{O}(Tn\log n)$ with additional computation for gradient and update. When we set $T=L$ with $L$ as the number of projections in SW, Max-SW is slower than SW. We ran Max-SW with $T=10$ in the gradient flow application, and we got this result.

We can see that Max-SW can make the gradient flow reduce the $W_2$ distance very fast in terms of the number of iterations. However, when compared in terms of real-time, it is much slower than SW (in Table 1). In addition, the flow from Max-SW cannot converge well (the $W_2$ score at the last iteration is higher than from SW) since it only uses one direction and can be seen as a minimax problem.

For PW, the orthogonal constraint is enforced via the Gram-Smith orthogonality process (or QR decomposition) which has the computational complexity is $\mathcal{O}(L^2 d)$ for $L$ is the number of projections (quadratic). Moreover, the sorting operator is more optimized in Python than Gram-Smith orthogonality process in Pytorch to the best of our knowledge. Therefore, PW is slower than SW. Hence, for the same computational time, PW has a smaller number of projections than SW which might lead to worse performance. We ran $L=3$ for PW, and got this result:

We see that the performance of PW is not as good as SW in Table 1 in both computation and quantitative scores.

Overall, we believe the comparative performance of distances depends on the applications i.e., each distance has its preferred applications. In our work, we focus on SW only i.e., focusing on the expectation with respect to the uniform slicing distribution. Since constructing control variates for SW is a new area of research, we believe the extension to other variants of SW is worth a careful future investigation.

We are happy to discuss more if the reviewer is not satisfied with our response.

Best regards,

We would like to express our gratitude for the time and feedback from the reviewer. We would like to extend the discussion as follows:

Q10:In Definition 1, the expectation of the controlled projected one-dimensional Wasserstein distance is non-nonegative (unbiased estimator of SWD). But what about the variable itself, $Z(\theta;\mu,\nu)$ , is it also non-negative?

A10: Thank you for your insightful question. The controlled variable $Z(\theta;\mu,\nu)$ can be negative, however, $\mathbb{E}[Z(\theta;\mu,\nu)]$ is still non-negative since it is the SW distance. Moreover, $\mathbb{E}[C(\theta)-B] =0$ which makes its estimators with Monte Carlo samples might be closed to 0 with high probability. Therefore, the Monte Carlo estimators of $\mathbb{E}[Z(\theta;\mu,\nu)]$ might be non-negative with high probability.

Q11: I'm wondering about the utility of Proposition 4, which states the left control estimator is equivalent to considering a control variate with respect to the 2-Wasserstein distance between a projection of Gaussian approximations of the origin measures. I'm thinking about the following scheme: we first calculate these Gaussian approximations of the origin probability measures then an SWD between these approximations and the controlled SWD is given by: |$\hat{SWD}$(origin measures)−$\hat{SWD}$(Gaussian Approximation of the origin measures)| without adding the factor $\gamma$ . I don't know it may highlight the importance of $\gamma$ to get more tight control variate. Any comment(s) will be appreciated for this point.

A11: First we would like to recall that SW-2 distance only has a closed-form when two measures are location-scale measures. For multivariate Gaussian with a diagonal covariance matrix and full covariance matrix, SW-2 is still approximated by the MC estimator.

In terms of approximation, approximating a high-dimensional measure by a Gaussian is likely to create more error than approximating an uni-dimensional measure by a Gaussian. Moreover, according to some central limit theorem [R1], the projected one-dimensional measure converges in distribution to a Gaussian measure. Therefore, we believe it is more appealing to do an approximation after the projection.

For the importance of $\gamma^\star$, we assume that you are referring to the difference estimator. In more detail, the estimator with $\gamma^\star$ is called the regression estimator. When $\gamma=1$, the corresponding estimator is named the difference estimator. While the difference estimator is simpler and does not need to deal with estimating $\gamma^\star$, the performance of the difference estimator depends strongly on the approximation error between the control variate and the interested integrand. For a bad approximation, the resulting difference estimator can lead to worse variance. In contrast, the regression estimator always leads to lower variance with a good approximation of $\gamma^\star$ i.e., can be guaranteed by increasing the number of MC samples. We would like to refer to Page 29, Chapter 8 in [R2] for a more detailed discussion.

[R1] Typical distributions of linear functionals in finite dimensional spaces of high dimension.

[R2] Monte Carlo theory, methods and examples

We appreciate the time and feedback from the reviewer. We would like to address the concerns from the reviewer as follow:

Q6: A major confusion is proposition 1, in which the paper claims to minimize KL divergence between discrete distribution and a Gaussian distribution. However, to the best of the reviewer's knowledge, there's no way to obtain finite value for KL between discrete distribution and continuous distribution, even under the most generalized setting with Radon–Nikodym derivative. The reviewer acknowledge that this potential flaw does not defeat the purpose of identifying the Gaussian proxy using information of $\mu$ and $\nu$(e.g. as an alternative, through moment matching), but urge the authors to revise this part accordingly.

A6: Thank you for your insightful comment. In the paper, we define KL divergence using the convention that $0\log 0=0$ i.e., $\lim_{x\to 0}x \log p(x) = 0$ as defining entropy in information theory literature. As commented by the reviewer, the main goal of defining the optimization problem is to find a Gaussian approximation of two discrete measures i.e., using the average of supports as the mean, and using the average square-L2 between supports and the mean as the variance. The formulation can also interpreted via Moment matching as suggested. We would like to recall that the Gaussian approximation can use any approximation schemes and the approximation in the paper is only a well-known type of approximation.

We have added this discussion to the paper in Section 3 in blue color. If the reviewer is not satisfied with the explanation with convention $0\log 0=0$, we will directly state the approximation formula instead of defining the optimization problem with the KL divergence.

[R1] Information, Physics and Computation, Chapter 1, page 3.

Q7:The usage of 'upper/lower bounds' seems to lack justification, as they are quite different from the Gaussian approach the paper proposes, in terms of how much correlation is changed/lost. In addition, though $\mathbb{E}_{\sigma_1(\theta;\mu)\sigma_2(\theta;\nu)}$ does not have a closed form, it's unclear why it can't be estimated by MC, as is done throughout the paper. (A personal thought: from a statistical perspective, in order to introduce correlation it might be possible to use $\mathbb{E}[\sigma_1^2(\theta;\mu)\sigma_2^2(\theta;\nu)]$ , if it is at all possible to address.)

A7: Thank you for your insightful comments. The main goal of the paper is to create a first attempt to connect the control variate literature and the sliced Wasserstein literature. To use control variate, we need to have a measurable function that has a tractable expectation and has a correlation with the original function. Motivated by the well-known result of the closed-form solution of Wasserstein distance between two Gaussians, we construct the control variate in the presented way. We would like to recall that the proposed control variates might not be optimal in the sense of reducing variance. They are designed as natural and computationally efficient control variates.

The reason why we do not estimate $\mathbb{E}[\sigma_1(\theta;\mu),\sigma_2(\theta;\nu)]$ via MC methods is because we do not want to create nested MC problems i.e., the expectation of the control variate is only intractable. In principle, we could use MC to estimate such quantity. However, it might lead to a little conflict with the original motivation and the conventional framework of using control variate. We have run an additional experiment in gradient flow with the control variate including estimating $\mathbb{E}[\sigma_1(\theta;\mu),\sigma_2(\theta;\nu)]$ via MC method. We have the following result:

Compared to Table 1 in the paper, despite being slightly better than SW without a control variate, this estimator is worse than the proposed control variates in the paper. As discussed, it might be due to the nested MC problem.

For using $\mathbb{E}[\sigma_1^2(\theta;\mu),\sigma_2^2(\theta;\nu)]$, to the best of our knowledge, the expectation is also intractable since it leads to a polynomial function of $\theta$ with the degree 4. Overall, we believe that the construction of the control variates is worth further investigation.

We would like to thank the reviewer for the time and feedback. We would like to answer questions from the reviewer as follow:

Q1: It is good that the authors provided strong motivations for using control covariates in Section 3.1, paragraph 2, the wiring in this part is poor. I suggest the authors follow the writing of [A. Shapiro 2021, Section 5.5.2] to re-write this part.

A1: Thank you for pointing out. We realized that there are two definitions of the controlled variable $Z(\theta;\mu,\nu)$ i.e., one with $\gamma$ and one with optimal $\gamma$. We have rewritten Section 3.1 paragraph 2 in blue color based on Section 5.5.2 in [R1]. We would be grateful if the reviewer could give comments on the rewritten paragraph.

[R1] Lectures on stochastic programming: modeling and theory

Q2: For the paragraph "Constructing Control Variates" in Section 3.2, the authors should omit the detailed deviation of the closed-form solution of $\text{W}_2^2(\mathcal{N}(m_1(\theta;\mu),\sigma_1^2(\theta;\mu)),\mathcal{N}(m_2(\theta;\nu),\sigma_2^2(\theta;\nu)))$. It is fine to present only the final simplified expression.

A2: Thank you for your suggestion. We have rewritten Section 3.2 based on your suggestion in blue color.

Q3: For the paragraph below Proposition 1, the description for computing the sliced Wasserstein between continuous distributions is brief and confusing. The authors may consider describing the algorithms in detail and present this part in an extra Appendix instead.

A3: Thank you for the suggestion. We would like to extend the discussion as follows:

Laplace Approximation. We are interested in approximating a continuous distribution $\mu$ by an Gaussian $\mathcal{N}(m,\sigma^2)$. We assume that we know the pdf of $\mu$ which is doubly differentiable and referred to as $f(x)$. We first find $m$ such that $f'(m)=0$. After, we use the second-order Taylor expansion of $\log f(x)$ around $m$:

$$\log f(x)\approx \log f(m) - \frac{1}{2\sigma^2}(x-m)^2,$$

where the first order does not appear since $f'(m)=0$ and \sigma^2 = \frac{-1}{\frac{d^2}{dx^2} \log f(x) \big{|}_{x=m}}.

Sampling-based approximation. We are interested in approximate a continuous distribution $\mu$ by an Gaussian $\mathcal{N}(m,\sigma^2)$. Here, we assume that we do not know the pdf of $\mu$, however, we can sample from $X_1,\ldots, X_n \overset{i.i.d}{\sim}\mu$. Therefore, we can approximate $\mu$ by $\mathcal{N}(m,\sigma^2)$ with $m=\frac{1}{n}\sum_{i=1}^n X_i$ and $\sigma^2 = \frac{1}{n}\sum_{i=1}^n (X_i-m)^2$ which is equivalent to doing maximum likelihood estimate, moment matching, and doing optimization in Equation 11 with the proxy empirical measure $\mu_n=\frac{1}{n}\sum_{i=1}^n \delta_{X_i}$. }

We have added the discussion in Appendix A in the revision in blue. We would be grateful if the reviewer could give us some comments.

Q4: I feel the theoretical analysis of the variance-reduced estimator is not enough. But when I look at related literature, there is little theoretical guarantees on this part. So I think it is fine regarding the theoretical contribution part.

A4: Thank you for your detailed comment. For a tractable form of the optimal coefficient $\gamma$ i.e., $Z(\theta;\mu,\nu,C(\theta)) in Definition 1, using control variates always leads to a lower variance (or at least the same), then lead to a lower approximation error in L1 norm. Nevertheless, the optimal coefficient is often intractable, hence, Monte Carlo samples are used to approximate the optimal coefficients. By using Monte Carlo samples, it is challenging to investigate the approximation error of the empirical estimator since it involves the ratio and the product of estimators of SW, covariance, and variances. As mentioned by the reviewer, the theoretical understanding of the approximation error is still an open area in the control variate literature even with some simple forms of the integrand. In our case, the integrand is SW which is a function on the unit-hypersphere (a manifold) and involves an optimal transport problem. Therefore, we would like to leave this investigation to the future work. Overall, we believe control variate can be mainly seen as a practical method to improve the performance of estimators in practice which has been shown in the paper with various applications.

We would like to thank you again for your time and constructive feedback!

Best regards,