Statistical and Computational Guarantees of Kernel Max-Sliced Wasserstein Distances

We thank the reviewer's positive comments and provide our response below:

• [relationship and distinctions with literature [9]?] We appreciate the reviewer for highlighting this point. Literature [9] establishes the statistical convergence rate for the empirical Max-Sliced (MS) distance as $O(R \cdot n^{-1/(2p)})$, where $R$ denotes the diameter of the sample space and $n$ is the sample size. This rate is minimax optimal, and the compactness assumption on the sample space is crucial to the analysis. In contrast, our results show that under the bounded kernel assumption (Assumption 1), the empirical KMS distance achieves the same convergence rate of $O(n^{-1/(2p)})$, without requiring the compactness of the original sample space. This allows our result to apply to a broader class of probability distributions. Furthermore, we prove that this rate is also minimax optimal.
  
  Our statistical analysis builds on the insights from [9]. Specifically, the KMS distance can be interpreted as first mapping the data distributions into an infinite-dimensional Hilbert space via an implicit feature map (induced by the kernel), and then applying the MS distance in that space (see our discussion in Remark 2.6). A key step in our proof leverages the statistical results for the MS distance in Hilbert spaces from [9, Corollary 2.8]. Our bounded kernel assumption ensures that the transformed distributions have bounded support (i.e., finite diameter) in the Hilbert space. We will include this discussion in our revision.

• [discussion of decreasing power issue of KMS?] We appreciate the reviewer for this insightful comment. Under fair alternative, we believe that the power of our KMS distance will decrease, as the strength of our assumption $\mathrm{KMS}(\mu,\nu) - \Delta(n,\alpha)>0$ increases. We will emphasize this point in our revision and add an additional experiment to illustrate the trend of decreasing power of our method compared to other baselines under such settings.

• [minimax lower bound?] We provide the following example to demonstrate that the bound $\Delta(n, \alpha)$ is tight. Consider the case where $\mathcal{B} = [0,1]$, the kernel is $k(x,y) = xy$, and the distribution is $\mu = \frac{1}{2}\delta_0 + \frac{1}{2}\delta_1$. Then the empirical KMS distance is given by $\mathrm{KMS}(\mu, \hat{\mu}_n) = \left|\frac{1}{2} - \frac{N}{n}\right|^{-1/p}$, where $N \sim \mathrm{Binom}(n, \frac{1}{2})$. It can be shown that $\mathbb{E}[\mathrm{KMS}(\mu, \hat{\mu}_n)] = \Theta(n^{-1/(2p)})$, thereby confirming that the rate $O(n^{-1/(2p)})$ is optimal in the worst case.

• [Motivation for rank analysis?] Our motivation for studying the rank of the solution is that it has impact on the quality of the resulting approximation to the original non-convex optimization problem. Recall that globally solving the KMS Wassertein distance involves a rank-1 constraint, which our SDR formulation relaxes it. When the solution to SDR is of rank higher than $1$, we resort to constructing a feasible rank-1 solution by taking its leading eigenvector. However, if the SDR solution is already low-rank, then the rounding procedure yields a solution that is closer to the ground-truth rank-1 optimum.
  
  We include an experimental study (see Figure 6 from the Anonymous link https://gofile.io/d/3Pcdxg) that visualizes how our rank reduction algorithm gradually reduces the rank of the SDR solution. Originally the optimal solution to SDR is of rank $400$, and our rank reduction algorithm iteratively reduces its rank by $1$ until we obtain the $19$-rank solution. We show the quality of the rounded feasible solution to the original KMS Wasserstein distance problem, and observe the corresponding objective will increase from $0.63$ to $0.68$.

• [Typos?] We appreciate the reviewer for pointing out our typos. We will correct them in the revision.

We very much appreciate reviewer's insightful comments and now provide our response on a point-by-point basis:

• [OT in RKHS?] Our formulation is closely related to Zhang et al. (2020), which considers Wasserstein distances between pushforward measures $\Phi(\mu)$ and $\Phi(\nu)$ via an implicit kernel map $\Phi$. In Remark 2.6, we show that our KMS distance is equivalent to Max-Sliced Wasserstein in this setting. A key distinction is that our method enjoys sharp statistical convergence rates, which are difficult to obtain in Zhang et al.’s framework due to the curse of dimensionality of Wasserstein. Moreover, while their work focuses on kernel embeddings, our motivation differs slightly and is to integrate nonlinear dimensionality reduction with OT. We will clarify this in the revision.

• [OT with a learnable nonlinear map?] We compare our method with $\mathcal{SW}(\Phi(\mu), \Phi(\nu))$ where $\Phi$ is a neural network. Such $\Phi$ is finite-dimensional and sensitive to neural network architectural choices. In contrast, our KMS can also be reformulated as the similar expression (in Remark 2.6) with $\Phi$ being an infinite-dimensional, non-parametric kernel mapping. In addition, our method has theoretical guarantee that $\mathcal{KMS}(\mu,\nu)=0$ iff $\mu=\nu$. This is important for many applications like two-sample testing, whereas it is not guaranteed in fixed neural network mappings.
  
  While we acknowledge that KMS has higher computational cost, it provides a complementary to neural networks. We will incorporate this discussion into the revision and hope the reviewer sees the merits of our proposed framework.

• [Generative Modeling justification?] We revised the setup following [Sliced-Wasserstein Autoencoder] with the new formulation

where $\psi$ and $\phi$ denote the decoder and encoder, respectively, and $q_{\text{prior}}$ denotes the pre-defined prior distribution on the latent space (uniform distribution on unit circle). We report experiment results in https://gofile.io/d/3Pcdxg (Table 3). Our KMS Wasserstein distance has competitive performance as indicated by the smallest Fréchet inception distance (FID) score. Also, it has learnt meaningful latent representation, possibly because it utilizes the flexible kernel-projection mapping to compare data distributions, and thereby achieve competitive performance in generative modeling.

We also clarify that the experiment setup for this part follows from reference [Sliced-Wasserstein Autoencoder]. It could be possible to improve the performance of all baselines by tuning neural network architectures, optimizers, training time, or even random seed. However, our focus in this paper is theoretical and not to develop state-of-the-art generative models. We hope this experiment will be enough to support the empirical value of KMS.

• [Complexity across all experiments?] The runtime for all methods is reported in the anonymous link (Table 2). While KMS is more expensive, the overhead is not prohibitive.

• [Time Complexity Concerns?] We will correct the stated complexity for solving (11) and add the reference as suggested.

• [Complexity of Algorithm 1?] The reviewer is correct. Time complexity of our Algorithm 1 is $\tilde{O}(n^3\delta^{-3})$, not $\tilde{O}(n^2\delta^{-3})$ as previously stated. In our initial submission, we followed established literature on first-order methods (e.g., [Nemirovsky and Yudin, Problem Complexity and Method Efficiency in Optimization]) which analyzes the complexity of constructing supgradient estimators at each iteration. However, we omitted the cost of the proximal gradient projection step induced by $h(S)$ in Eq. (13). This step involves computing the matrix exponential and matrix logarithm, each of which requires $O(n^3)$ time. Fortunately, these operations can be efficiently implemented using well-established software packages.

• [Complexity of solving SDR and rank-reduction?] Solving KMS-Wasserstein is NP-hard; we only analyze its convex relaxation. We also refer to our response to Reviewer LEjS for motivation on the rank-reduction algorithm. We acknowledge that further calibrating the optimal solution from SDR to low-rank space is computationally expansive (whose complexity is $\mathcal{O}(n^5)$), but it is still of theoretical interest if we wish to benefit from the superior performance induced by low-rankness of solution.

• [Typos?] We thank the reviewer and will correct all noted typos and clarify notations in the revision.

We very much appreciate reviewer's positive comments and now provide our response on a point-by-point basis:

• [Complexity of KMS?] We would like to clarify that the time complexity of our Algorithm 1 is $\tilde{O}(n^3\delta^{-3})$, not $\tilde{O}(n^2\delta^{-3})$ as previously stated. In our initial submission, we followed established literature on first-order methods (e.g., [Nemirovsky and Yudin, Problem Complexity and Method Efficiency in Optimization]) which analyzes the complexity of constructing supradient estimators at each iteration. However, we omitted the cost of the proximal gradient projection step in Eq. (13). This step involves computing the matrix exponential and matrix logarithm, each of which requires $O(n^3)$ time. These operations are efficiently implemented in modern linear algebra libraries, but their asymptotic cost remains cubic in $n$.
  
  Moreover, although Algorithm 1's complexity is independent of the data dimension $d$, it does require precomputing the Gram kernel matrix $G$ (Eq. (8)), and generating the vectors $M_{i,j}'$ as input. This preprocessing step incurs a time complexity of $O(n^3d)$, where the factor of $d$ arises from the kernel computation $k(x,y) = \exp(-\|\|x - y\|\|_2^2 / \sigma^2)$, which scales linearly with $d$ for $x, y \in \mathbb{R}^d$.
  
  In summary, computing the KMS distance requires time complexity $O(n^3(\delta^{-3} + d))$, which is significantly larger than that of the vanilla sliced Wasserstein distance. We will include a detailed discussion of this point in our revised manuscript.

• [Complexity and Performance Trade-off?] While KMS has in general the highest computational time among the evaluated methods, the overhead is not prohibitive. Importantly, it consistently achieves superior performance, demonstrating that our method is well-suited for practical machine learning applications, albeit with increased computational demands.
  
  Please also see Table 2 from the Anonymous link https://gofile.io/d/3Pcdxg that reports the numerical running time for hypothesis testing, change-detection, and generative modeling experiments. We observe for the change-detection experiment, our approach even has smaller computational time compared with Sinkhorn Divergence, SW, and GSW. This efficiency arises because the nonlinear projector can be precomputed using pilot data, enabling fast online computation of the test statistic. In contrast, the baseline methods must recompute the statistics at each detection step, resulting in longer runtimes. We will make this point explicit in the revision to highlight the practical utility of our method.

• [Typos?] We appreciate the reviewer for pointing out our typos. We will correct them in the revision.