Fourier Sliced-Wasserstein Embedding for Multisets and Measures

Response to summary:

We would like to highlight that, in addition to the theoretical guarantees for our embedding, we present the impossibility result stated in Theorem 4.4. This result proves that it is impossible to embed discrete distributions into any finite-dimensional Euclidean space in a bi-Lipschitz manner. This saves the community further efforts in that direction, and essesntially shows that an embedding with substantially better analytical properties than the FSW does not exist.

Weakness 1: Introducing baseling methods

Thank you for your suggestion. We added to the revised manuscript a paragraph describing all the baseline methods (lines 498-504).

We would like to stress that the methods presented in Table 2 are those introduced in [Chen] and earlier papers, and were not implemented by us. Specifically: $\mathcal{N}\_{\textup{ProductNet}}$, $\mathcal{N}\_{\textup{SDeepSets}}$ and WPCE use their own architectures, and Sinkhorn is an approximation algorithm specifically designed to approximate the $p$-Wasserstein distance.

Our architecture based on $E_1$, $E_2$, $\Phi$, described in line 485, achieves state of the art results with a simple combination of our embedding and one MLP. In comparison, $\mathcal{N}{\_\textup{ProductNet}}$ produces inferior results using three MLPs.

The only method other than ours that we tested with our architecture was PSWE, which is designed to compute Sliced-Wasserstein preserving embeddings.

Lastly, the reason why Sinkhorn cannot be incorporated into our architecture is due to its own inherent limitation: it takes pairs of input distributions and estimates their distances, rather than computing a distance-preserving embedding for individual distributions. This significantly limits its applicability to practical learning tasks, as we further discuss in our response to Reviewer Ew1o. We also added a brief explanation in line 307.

Weakness 2: Real-data application of our method

The experiment presented in Table 2 illustrates the utility of our embedding in a learning task on real-world data (ModelNet-40). We note that our paper includes all experiments from the NeurIPS-accepted work by Chen and Wang [Chen], as well as theoretical results that in our opinion merit acceptance in their own right.

To be continued...

Weakness 3: Why bi-Lipschitzness is important

The lack of bi-Lipschitzness, which plagues most prevalent multiset architectures to date, practically implies that there inevitably exist pairs of different input multisets that appear numerically identical to the architecture. This poses a problem, for example, in classification tasks where such pairs need to be assigned different labels. Any multiset architecture based on sum- or max-pooling is provably affected by this problem [FWT].

Achieving a bi-Lipschitz embedding for multisets has been recognized as an important goal in previous works. Our work is the first to fully achieve this goal. Below is a selection of previous works that underscore the importance of bi-Lipschitzness:

"The question of which metric spaces admit a bilipschitz embedding into some (finite-dimensional) Euclidean space is natural, and has received a lot of attention in recent work. The results obtained so far indicate that there is no simple answer to this question."

— Lang, Urs, and Conrad Plaut. "Bilipschitz embeddings of metric spaces into space forms." Geometriae Dedicata 87 (2001): 285-307.

"Since the late 1990’s, it has become apparent that designing efficient approximate nearest neighbor algorithms, at least for high-dimensional data, is closely related to the task of designing low-distortion embeddings. A bi-Lipschitz embedding between two metric spaces $(X,d\_X)$ $(X',d\_{X'})$ is a mapping $f : X \to X'$ such that ... where the parameter $D \geq 1$ called [sic] the distortion of $f$."

— Indyk, Piotr, and Assaf Naor. "Nearest-neighbor-preserving embeddings." ACM Transactions on Algorithms (TALG) 3.3 (2007): 31-es.

"The second negative result is that while moments of MLPs with analytic activations can be injective, they can never be stable in the bi-Lipschitz sense. This points to a possible advantage of injective multiset functions that are not based on moments, but rather on sorting or max-filters."

— Amir, T., Gortler, S., Avni, I., Ravina, R., & Dym, N. (2023). "Neural injective functions for multisets, measures and graphs via a finite witness theorem." Advances in Neural Information Processing Systems (NeurIPS) 37 (2023)

“We propose developing fine-grained expressivity results, namely metric equivalencies between explicit graph metrics and feature metrics for GNNs on graphs with features. An ideal result would derive a bi-Lipschitz correspondence between such metrics.”

— Christopher Morris, Nadav Dym, Haggai Maron, İsmail İlkan Ceylan, Fabrizio Frasca, Ron Levie, Derek Lim, Michael Bronstein, Martin Grohe, and Stefanie Jegelka. "Future Directions in Foundations of Graph Machine Learning." International Conference on Machine Learning (ICML) (2024)

Question 1: Regarding point (2), is "Multisets" simply another term for "discrete distributions"?

Multisets and discrete distributions refer to different concepts. Multisets are essentially sets that account for repetitions. For instance, $\\{b,a,b\\} = \\{a,b,b\\} \neq \\{a,b\\}$. Discrete distributions, on the other hand, are probability distributions with finite support. However, multisets can be idenitified with the subset of of discrete distributions with uniform weights, as discussed beginning at line 183.

Question 2: Could you clarify what "E2" refers to in lines 473-474?

$E_1$ and $E_2$ are two independent instances of the FSW embedding, with different input and output dimensions. $E_1$ maps distributions over $\mathbb{R}^d$ into $\mathbb{R}^{m_1}$, and $E_2$ maps distributions over $\mathbb{R}^{m_1}$ into $\mathbb{R}^{m_2}$, with $m_1$, $m_2$ being architecture hyperparameters. This is detailed in the manuscript (lines 481-485 in the revised version).

References:

[Chen] Chen, S., Wang, Y. "Neural approximation of Wasserstein distance via a universal architecture for symmetric and factorwise group invariant functions." Advances in Neural Information Processing Systems (NeurIPS) 37 (2023).

[FWT] Amir, T., Gortler, S., Avni, I., Ravina, R., & Dym, N. (2023). "Neural injective functions for multisets, measures and graphs via a finite witness theorem." Advances in Neural Information Processing Systems (NeurIPS) 37 (2023).

Thank you :)

We thank the reviewer for the comments. We have uploaded a revised version of the manuscript, where the comments have been addressed and incorporated.

Response to summary:

We would like to highlight an additional theoretical contribution in our paper: the impossibility result of Theorem 4.4, which proves that it is impossible to embed discrete distributions into any finite-dimensional Euclidean space in a bi-Lipschitz manner. This saves the community further efforts in that direction, and essesntially shows that an embedding with substantially better analytical properties than the FSW does not exist.

Response to weaknesses:

1. On the inclusion of $\mathcal{W}_{\infty}$: We included the definition of the $p$-Wasserstein distance with $p = \infty$ to allow our results to be stated across all $p \in [0, \infty]$. Our bi-Lipschitzness guarantee and impossibility result apply uniformly for all $p$ in this range.
2. Complexity of Wasserstein when $d=1$: The complexity in this case is the computational complexity of the sort function, which is $\mathcal{O}(n \log n)$. This is stated in lines 223-224.
3. Definition of STD: STD here is the Standard Deviation. We clarified this in the revised manuscript, l. 341-342.
4. Proof ideas in the main text: We added an overview of the proof ideas of Theorems 4.1, 4.2 and 4.4 to the revised manuscript. Thank you for this comment.

Response to question on the practical motivation:

To illustrate the advantage of our approach for practical applications, consider a learning task on multisets handled by traditional architectures based on sum- or max-pooling. With these methods, certain pairs of input multisets may appear numerically identical, meaning the architecture will not be able to distinguish between them—even if they represent different underlying data. In contrast, our approach, due to its bi-Lipschitzness guarantee, can distinguish between these multisets in a way that reflects their actual differences. This practical advantage is evidenced in our experimental results shown in Table 2.

There is a fundamental difference between our embedding approach and approaches such as Distributional Sliced Wasserstein and Max-Sliced Wasserstein. Our approach takes one input distribution at a time and computes an embedding, whereas the aforementioned approaches take two input distributions and estimate a distance between them. Pairwise methods have two disadvantages in comparison with embeddings: (i) higher computational complexity when computing multiple pairwise disances, and (ii) limited applicability to real-world learning problems.

In terms of applicability, a pairwise method cannot be directly applied to common learning tasks, such as object classification, where the inputs are typically individual distributions. In contrast, an embedding is readily applicable to such tasks, as demonstrated in our experiments.

In terms of computation, pairwise methods to estimate sliced optimal transport distances typically require $\tilde{\mathcal{O}}(mnd)$ time (neglecting logarithmic factors), where $m$ is the number of slices, $n$ is the maximal number of support points, and $d$ is the ambient dimension of the support. Thus, computing all pairwise distances for a set of $k$ distributions would take $\tilde{\mathcal{O}}(k^2 mnd)$ time. In contrast, computing our embedding takes $\tilde{\mathcal{O}}(mnd)$ time for each input distribution, and pairwise distances can then be computed in the Euclidean space $\mathbb{R}^m$, resulting in a considerably lower total complexity of $\tilde{\mathcal{O}}(k mnd + k^2 m)$. This approach is therefore significantly more scalable for large datasets where pairwise distance computations are required.

We appreciate the reviewer's comment and will clarify this in the paper.