Fourier Sliced-Wasserstein Embedding for Multisets and Measures
Authors: Tal Amir
Nadav Dym
Venue: ICLR 2025 Poster (2025)
Keywords: Sliced Wasserstein distance
Euclidean embedding
multiset embedding
bi-Lipschitz
permutation invariant
Primary area: learning on graphs and other geometries & topologies
TL;DR: A bi-Lipschitz Euclidean embedding for multisets that extends injectively to measures, based on the sliced Wasserstein distance

Abstract: We present the _Fourier Sliced-Wasserstein (FSW) embedding_—a novel method to embed multisets and measures over $\mathbb{R}^d$ into Euclidean space.

Our proposed embedding approximately preserves the sliced Wasserstein distance on distributions, thereby yielding geometrically meaningful representations that better capture the structure of the input. Moreover, it is injective on measures and _bi-Lipschitz_ on multisets—a significant advantage over prevalent methods based on sum- or max-pooling, which are provably not bi-Lipschitz, and, in many cases, not even injective.
The required output dimension for these guarantees is near-optimal: roughly $2 N d$, where $N$ is the maximal input multiset size.

Furthermore, we prove that it is _impossible_ to embed distributions over $\mathbb{R}^d$ into Euclidean space in a bi-Lipschitz manner. Thus, the metric properties of our embedding are, in a sense, the best possible.

Through numerical experiments, we demonstrate that our method yields superior multiset representations that improve performance in practical learning tasks. Specifically, we show that (a) a simple combination of the FSW embedding with an MLP achieves state-of-the-art performance in learning the (non-sliced) Wasserstein distance; and (b) replacing max-pooling with the FSW embedding makes PointNet significantly more robust to parameter reduction, with only minor performance degradation even after a 40-fold reduction.

Published as a conference paper at ICLR 2025
FOURIER ELICED-WASSERE EMBEDDING
FOR MULTISETS AND MEASURES
Tal Amir
Nadav Dym
Faculty of Mathematics
Faculty of Mathematics and Faculty of Computer Science Technion-Israu Institute of Technology Haifa, Israel
Technion-Isracl Institute of Technology Haifa, Israel
t lami remirrtechnion .n.ichhniin.eee. acc. 11
nadavdym@techh on ac
ABSTRACT
We present the Fourier Sliced- Wasserstein (FSW) embedding novel method embed multisets and measures over IR into Euclidean space.
Our proposed embedding approximately preserves the sliced Wasserstein distance on distributions, thereby yielding geometrically meaningful representations that better capture the structure of the input. Moreover, it is injective measures and bi-Lipschitz on multisets-a significant advantage over prevalent methods based sum- max-pooling, which are provably not bi-Lipschitz, and, in many cases, not even injective. The required output dimension for these guarantees is near- optimal: roughly 2Ndd, where N is the maximal input multiset size.
Furthermore, we prove that it is impossible to embed distributions over Rd into Euclidean space in a bi-Lipschitz manner. Thus, the metric properties of our em- bedding are, in a sense, the best possible.
Through numerical experiments, we demonstrate that our method yields supe- rior multiset representations that improve performance in practical learning tasks. Specifically, we show that (a) a simple combination of the FSW embedding with MLP achieves state-of-the-art performance in learning the (non-sliced) Wasser- stein distance; and (b) replacing max-pooling with the FSW embedding makes PointNet significantly more robust parameter reduction, with only minor per- formance degradation even after a 40-fold reduction."
INTRODUCTION
Multisets are unordered collections of vectors that account for repetitions. They are the main math- ematical tool for representing unordered data, with perhaps the most notable example being point clouds. As such, there is growing interest in developing architectures suited for learning tasks on multisets. To address this need, several permutation-invant duced, with applications for point-cloud classification (Qi et al., 2017a), chemical property pre- diction (Pozdnyakov Ceriotti, 2023), and image deblurring (Aittala & Durand, 2018). Multiset aggregation functions also serve as key components in more complex architectures, such Mes- sage Passing Neural Networks (MPNNs) for graphs (Gilmer et al., 2017), or setups with multiple permutation actions (Maron al., 2020).
neural networks have been intro-
A central concept in the study of multiset functions, i.e. functions defined on multisets, is injectivity. Its importance is highlighted by the following observation: An architecture that cannot separate two distinct multisets X * X' will not be able to approximate a target function f that distinguishes between these multisets, i.e. s(X) * I [(X).. Conversely, model that maps multisets injectively to vectors, composed with an MIP, can universally approximate all continuous multiset functions (Zaheer et al., 2017: Dym & Gortler, 2024). This observation has inspired works study the injectivity of multiset architectures (Wagstaff et al., 2022; 2019; Tabaghi & Wang, 2024). Injectivity on multisets also plays a key role in the development of expressive MPNNs (Xu et al., 2018).
Our code is available at https:////// cow.b.coo/tatt r/FSW embedding.

---

Published as a conference paper at ICLR 2025
Common multiset architectures typically based on simple building blocks of the form
E (f21 x, = Pool F (zx1)) ((((")....F((())) ((.()))
where F is usually MLP, and Pool is simple pooling operation such maximum, sum, or mean. Xu et al. (2018) showed that multiset functions based on max- or mean-pooling are never injective, but injectivity can be achieved using sum-pooling, under the assumption that the vectors (:()) come from a discrete domain, and an appropriate function F is used. Then it was shown by Zaheer et al. (2017); Maron et al. (2019) that injectivity over multisets with continuous elements can be achieved using sum pooling with a polynomial P. The more common case, in which F is neural network, was studied in (Amir et al., 2023). There was shown that injectivity on multisets and measures over Rd can be achieved using F that is shallow MLP with random parameters and analytic, non-polynomial activations, such as Sigmoid or Softplus.
However, injectivity alone is not the strongest property may desire for multiset functions. While an injective multiset embedding E can separate pairs of distinct multisets X * X', this does not ensure the quality of the separation. Ideally, if two multisets X, X' are far apart in terms of some notion of distance, then one would expect ( ), E'(X) E 1 to also be far apart, and vice versa. The standard mathematical notion used to guarantee this behaviour is bi- bi-Lipschizzess.
Definition, Let E: 'D - R." where D is collection of multisets or, more generally, distributions over Idd. We say that E' is bi-Lipschitz with respect to Wp, if there exist constants c, C such that
c. (., I [[E(p) E - C Wp(u, P), V4, DD
(1)
where W, denotes the p- Wasserstein distance and || denotes the norm.
The Wasserstein distance, defined in the next section, serves as a standard notion of distance for multisets and distributions. The ratio of the Lipschitz constants C/c in (1) represents a bound on the maximal distortion incurred by the map E, analogous to the condition number of a matrix.
Bi-Lipschitz embeddings enable the application of metric-based learning methods, such as nearest- neighbor search, data clustering, and multi-dimensional scaling, to non-Fuclidean data types like multisets and measures. These methods can be more readily applied to the embedded Euclidean domain than to the original data domain, where metric computations are typically more expensive (Indyk& & Thaper, 2003). The bi-Lipschitzness of the embedding often provides correctness guaran- tees for these applications, which depend on the Lipschitz constants and C (Cahill et al., 2024).
Achieving bi-Lipschiszeess is typically more challenging than injectivity and often requires different theoretical tools. Recently, in (Amir et al., 2023), we proved that methods based on average- or sum-pooling can never be bi-Lipschitz, while Xu et al. (2018) showed that methods based on max- pooling cannot even be injective. Thus, bi-Lipschitz embedding would constitute substantial improvement over the standard multiset embedding methods currently in use.
To date, two main approaches have been proposed for constructing bi-Lipschitz multiset embed- dings: (1) sort embedding (Balan et al., 2022), which is based on sorting random projections of the multiset elements, and (2) max filtering (Cahill et al., 2022), which is based on computations of Wasserstein distances from 'template multisets', and is computationally expensive.
While bost-based methods have been used with some success (Zhang al., 2019; 2018: Balan al., 2022), their popularity in practical applications remains limited, despite their bi- Di-Liipchitzeess guar- antees. A likely explanation for this is that these methods can only handle multisets of fixed size, and, to date, there is no known way to generalize them to multisets of varying size, let alone to dis- tributions. This is major limitation, since multisets of varying size arise naturally in many common learning tasks, for example graph classification, where vertices may have neighbourhoods of differ- ent sizes. This problem often circumvented via ad-hoe solutions such as padding (Zhang et al., 2018) or interpolation (Zhang al., 2019), which do not preserve the original theoretical guarantees of the method. Moreover, even in the restricted setting of fixed-size multisets, the bi-Lipschitzness guarantees of these methods typically require prohibitively high embedding dimensions.
Our goal in this paper is to overcome these limitations by constructing a bi- bi-Lipsehitz embedding for the collection of all multisets over dd with at most N elements. We denote this collection by Sgx (Rd). Note that the assumption of bounded cardinality is necessary, as otherwise, even injectivity is impossible; see (Amir et al., 2023, Theorem C.3). We also interested in the larger
2

---

Published as a conference paper at ICLR 2025
collection of probability distributions over R supported on most N points, which we denote by Pe (R*). This setting, in which the points may have non-uniform weights, can be particularly sets (Lee et al., 2019), well graph architectures such
relevant for attention-based methods as GCN (Kipf & Welling, 2016) or GAT' (Veličković et al., 2018), which use non-uniform weights for vertex neighhourhoods. In summary, our main goal is:
Main goal For D = SeN (R*) and D - PsN (R"), construct an embedding E : D + R" that is injective and, preferably, bi-Liihchitz.
Main results We propose an embedding method for multisets and distributions with finite support, which is a non-trivial generalization of the sort embedding. We observe that the Euclidean distance between the sort embeddings of two multisets can be interpreted their sliced Wasserstein distance (Bonneel et al., 2015): in the special case where the input consists of multisets of fixed size, this sampling corresponds to the project-and- anddctt----oorttooo operations used in the sort embedding. Based on this interpretation, we extend beyond fixed-size multisets and propose an embedding method for both Sex (TR)) and P&N (R4). Our method essentially operates follows: (1) compute random one-dimensional projections (also called slices) of the input distribution; (2) compute the quantile function of each projected distribution; and (3) sample each quantile function at random frequency in the Fourier domain. We name our method the Fourier Sliced- Sereed--Wsseetein (FSW) embedding and denote by EFSW The function EFSW
finite Monte Carlo sampling of
PPN (R²) IR is of the form
= FESW (n) = EESW (15 (vS(t) g(h)) k=
It maps multisets and distributions into R" and depends on parameters (k)) Rd, g({) E R, k [m], representing projection vectors and frequencies respectively. has the following properties:
1. Bi-Lipschitzness on multisets: For m 2 2N the map EFSW SeNN (R.) + R" is bi- Lipschitz (and, in particular, injective) for almost any choice of the embedding parameters km (Theorem 4.1 and Corollary 4.3).
(118), s c(k)) m
2. Injectivity on distributions and measures:
For m 2 2Nd + 2N - 1, the map EtSw PeN (R2) IRm is injective (but not bi-Lipschitz) for almost any choice of parameters (Theorem 4.... Moreover, by adding one more output coordinate, the embedding can be ex- tended from distributions to measures with arbitrary total mass while preserving injectivity; see Appendix A. Furthermore, we prove that it is impossible to construct bi-L.pschiiz Euclidean embedding for P&N (R&) (Theorem 4.4). This suggests that the metric proper- ties of our embedding are, in sense, the best possible.
3. -rseed--Wssrsttein approximation: The expectation of 1 FSW (n) - E [ over m the parameters (v(+), _(()) k= drawn from appropriately defined distribution, is ex- actly the squared sliced- Wasserstein distance between μ and j (Corollary 3.3), with the standard error decreasing as o( ViT
4. Differentiability: The map FFSS is continuous and piecewise smooth in both the input N
m measure parameters (ax(+), wi) 1= and the embedding parameters (v(*), c(()) k=1 (see e.g.
(25)). Thus, is amenable to gradient-based learning, and its parameters can be trained.
5. Complexity: The embedding EtS(()) can be computed efficiently in
O(mNd + mN log N) time-the same complexity (up to the logarithmic term) as the most efficient methods used in practice (Theorem A.2, Appendix A.2).
In Properties and above, the required embedding dimension m near-optimal, essentially up to multiplicative factor of two.
Empirically, in Section 5, we demonstrate the practical promise of our method by evaluating on the task of learning the (non-sliced) I-Wasserstein distance function. We show that a simple composition our embedding with an MLP outperforms the sfate-of-the-eee while requiring shorter training times. Moreover, significantly improves the architecture's 's robustness in the low- now-merameter regime, with only minor performance degradation even after 40-fold parameter reduction.
show that replacing max-pooling in PointNet with our embedding
3

---

Published as a conference paper at ICLR 2025
2
PROBLEM SETTING
In this section, we describe the problem detail and briefly review its theoretical background and existing approaches.
2.1 THEORETICAL BACKGROUND
We begin by defining the spaces of multisets and distributions of interest, along with the metrics we use on these spaces.
Multisets and distributions Denote by (S)( the collection of all probability distributions over a measurable domain f2 S IR,, supported on at most N points. Any distribution E Pgv can be parametrized by N points (1)) E Ω and weights 2 0 such that
N
w;d i=1
where w (W),.... w is in the probability simplex A R. and o Dirac's delta function at x. Distributions with fewer than N support points can be parameterized in this way by setting some of the weights w; to zero and choosing the corresponding (( arbitrarily.
Similarly, let (S) be the collection of all multisets over S Rd with most N ele- ments. To enable computation of Wasserstein distances, which are defined on distributions. we modify the definition of (I2( by excluding the empty multiset and identifying each multiset x = {x(0)" 1==1 E (2)( with the distribution x4[]] E Ps2(( that assigns uniform weights wy to all (f)), taking their multiplicities into account." Throughout this work, our use of SeN (52) assumes this modification, and we regard SgN as subset of PCN((2).
Note that embedding in its basic form, be defined in the next section, will take input multi- sets from SEN (12). As such, it will not defined on the empty multiset, and will not distinguish between multisets of different sizes if their element proportions are identical. This can be casily remedied by augmenting the embedding with additional output coordinate that encodes the mul- tiset cardinality, or in the case of measures, the total mass Σ wi; see Appendix A. for details.
Throughout this work, we focus Ω ded and only discuss finitely-supported multisets and distri- butions. Nonetheless, our embedding can accommodate general distributions over Rd, while retain- ing its Siiced-Wassersterr approximation property. Thus, principle, our method can be applied to structures other than point clouds, for example, polygonal meshes and volumetric data.
Wasserstein distance As distance measure on S&N (Rd)) and P&N ( we use the aforemen- tioned Wasserstein distance. Intuitively, this distance is the least amount of work required to 'trans- port' one distribution another. For two distributions μ. f E P&N (Re)), parametrized by points (e))) a() and weights wa, Vy as in (2), the p- Wasseesstnin distance between and jı is defined by
2c01pp)
())
inf FETT(AR))
Σ Thj
p E (0,
[JJE[N]
Woo (H,
inf =EN((uF))
max {1(000)
x60ه
i.j [N], mij > 0).
where || || is the Euclidean norm, and 11(H, ū is the set of all transport plans from to jt:
Mpp NERN (Vi.j (N]) Fif A
Σ ল.z mリ=
w.:
Taj Wij
JE[N]
SE[N]
2 For example, the multiset X (a,b, b) E Sggs is identificd with the distribution xx(x] = = 2 2 566
(R)( c.g. X b,a,,, and Y aa, b, are considered identical in Sge(R), since x(X] ,\
4

---

Published as a conference paper at ICLR 2025
The numbers Ty intuitively represent the amount of mass to be transported from point x0 to z()) Whenever p is omitted, we refer to Wp with = 2. Similarly, always denotes the &g norm.
Computation of Wasserstein The Wasserstein distance can be computed in ( log N) time via a linear program (Altschuler et al., 2017; Orlin, 1988). Alternatively, one may the Sinkhorn approximation algorithm (Cuturi, 2013), which takes 0(N2g time, being the tolerance (Altschuler et al., 2017). This was improved to Ô Omin { [ N2 in (Dvurechensky et al., 2018). However, in the special case d 1, it can be computed significantly faster.
Wasserstein in ID In one dimension, the Wasserstein distance can be computed only O(N log N) time. If (2....., .n), y (VI,..., Un) two vectors in R", then the dis- tance between the two uniform distributions p|æ)), ubyl over the vector coordinates is given by
W(ula)),
[sot( sort(y)ll,
with sort : 1n + R being the function that returns the input coordinates sorted in increasing order.
For arbitrary distributions in (R), the Wasserstein distance can be computed sim- ilarly via the quantile function. For a distribution I over R, the quantile func- tion Qμ 1) 7 R is continuous analogue of the sort function, defined by
(u) = inf (xr E m((-oo, l).
[4,5,6]
Figure/ depicts the quantile functions for three different multisets.
(3,4)
The quantile function enables an explicit formula for the Wasser- stein distance between two distributions over (see e.g. Bayraktar & Guo (2021), Eq. 2.3 and the paragraph thereafter):
W(n, V. (Qx(1) Qa(()
[1,2,3)
(4)
Note that when and are generated by multisets of the same cardinality (like the two multisets of cardinality 3 in Figure 1), the formulas (4) and (3) coincide.
1
Figure Quantile functions for three different multisets
Sliced Wasserstein distance The sliced Wasserstein distance, proposed by Bonneel et al. (2015) surrogate for the Wasserstein distance, exploits the efficient computation of the latter for to define a more computationally tractable distance for > 1. It is defined as the average Wasserstein distance between all oimeeimnnsional projections (or slices) of the input distributions. For formal definition, we first define the projection of distribution.
Definition. Let μ Σn= w;6min PEN (R.). The projection of in the direction e Rd. denoted by o" μ., is the one-dimensional distribution in P&v(R) defined by
47
11
Using the above definition, the sliced Wasserstein distance between μ. f PSN (R.) is defined by
SW(, = V EE.(M?(0 vT))),
(5)
where W is the squared 2- Wasserstein distance, and the expectation E,[ is over the direc- tion vector Uniform gd-1), 1), i.c. distributed uniformly over the unit sphere in gd
2.2 EXISTING EMBEDDING METHODS
We now return to our main goal of constructing an embedding E PSN (R*) + RT In this subsection, discuss existing embedding methods and some straightforward ideas extend them.
5

---

Published as a conference paper at ICLR 2025
We first observe that, on the collection of multisets over I of fixed size n, it follows from (3) that the map fz1,..., 2n)
sort(21...... ,2n) is an isometry, ie. (1) holds with C' = c=
To extend this idea to multisets in Sgv(RR, naive approach would be to represent each multiset in SeN( (I) by a multiset of size N. with N being the least common nultiple (LCM) of {1. 2.... NN. For example, for N = 3. LCM 2, 2,3}) = 6, and thus multisets in SgN (IR) of sizes 1 {a}. 2 {a, b). and fa, b, c), would be represented, respectively. by aa,a, a, a, a, a, a), {a, a, a, b, b, b) and {a, b, b, c, c). At this point, a sorting approach can be applied. However, as N increases, this method quickly becomes infeasible, both in terms of computation time as well memory, since N grows exponentially in N. Moreover, this method cannot handle arbitrary distributions in P&N (IR, whose weights may be irrational.
One possible approach to embed general distributions E PEN (R) is to sample (u(() at m points 4m E [0, 1] equispaced or drawn uniformly at random. While this approach would indeed approximately preserve the Wasserstein distance, as follows from (4), it is easy to show that for any finite number of samples m, this embedding is not injective on PeN (R). Moreover, it is discon- tinuous with respect to the sampling points fg and input probabilities wy, and thus not amenable to gradient-based learning methods. Our method, described in the next section, resolves these issues by sampling the quantile function in the frequency domain rather than in the f-domain.
When considering P& (R*) with > 1, one natural idea is to take m me-dimensional projections of the input distribution, and then embed each of the projections using one of the methods described above for PgN((R) In the case of multisets of fixed cardinality n, this corresponds to the mapping
{ai,.... Fn
Vmn rowsort [o]] a: ) bE(m),)e))))
This idea was discussed in (Balan et al., 2022; Zhang et al., 2019; Dym & Gortler, 2024; Balan & Tsoukanis, 2023b). In expectation over the directions 2%. this method indeed gives a good ap- proximation of the sliced Wasserstein distance. The relationship the d-dimensional Wasserstein distance is priori less clear. Balan & Tsoukanis (2023a) showed that for m that is exponential in n, this mapping bi-Lipschitz for almost any choice of v1...., Um- Later, Dym & Gortler (2024) showed that m = 2nd + is sufficient. In this paper, we combine this idea of using linear projec- tions, with our idea of Fourier sampling of quantile functions, to construct an embedding defined on the whole of P (RR) while maintaining theoretical guarantees and practical efficiency.
In a related line of work, Kolouri et al. (2015); Naderializadeh et al. (2021); Lu et al. (2024) devel- oped method that embeds continuous distributions into infinite- dimensionll Hilbert space iso- metrically with respect to the sliced Wasserstein distance. However, in practice, a finite-dimensiona. discretization is used, which breaks the method's 's injectivity, we show empirically in Section 5. In contrast, our method is provably injective with a finite and near-optimal embedding dimension of 2Nd.
Lastly, Haviv et al. (2024) recently proposed a transformer-basee architecture called Wasserstein Wormhole, designed to compute Euclidean embeddings for multisets and distributions. This archi- tecture can be trained to approximate the Wasserstein distance. However, its reliance on training is a key limitation, it does not guarantee Wasserstein distance preservation or even injectivity. This concern becomes particularly significant when generalizing to
samples.
3 PROPOSED METHOD
Our method for embedding a distribution f essentially consists of computing random slices o o and, for each slice, taking single random sample of its quantile function (t), However, instead of sampling the function directly, we sample its cosine transform-a variant of the Fourier transform. Since the Fourier transform is linear isometry, integrating the squared difference of these samples for two distributions μ. 인 yields the squared sliced Wasserstein distance SW ū)), as we shall show next. We will also show that this sampling method guarantees injectivity. unlike direct sampling of (t) Lastly, since Q,TT,(t) is compactly supported step function, its Fourier transform is smooth with respect to the frequencies, and thus so is our embedding. We now discuss this in detail.

---

Published as a conference paper at ICLR 2025
Definition 3.1. Given projection vector g4-1 and number & 2 0 denoting frequency, define the me-sammpe embbeddigg ETSW (1;0,5) PEN(RR) R by
EFS (k) 2(1 E) I. Qut cos (2mtt)dt,
(6)
which is the cosine transform of (t), sampled at frequency & and multiplied by + E; see Ap- pendix B.1 for further discussion. Details on the practical computation of ERSW are in Appendix A.2.
Next, we define probability distribution Dg for the frequency ç, given by the PDF
fc({) (1 (1+
ç 1
We now show that ESS preserves the sliced Wasserstein distance in expectation over v, Theorem 3.2. (Proof in Appendix B.2) Let j E P& (R)). with Bn C Rd being the closed fo-n ball of radius R, centered at zero. Let v0 ~ Uniform [g ) ~ Dz. Then
2 = SW (1,(i). FETS (15 (, FFS (Ø=; 0, 6)
(7)
[L (1, FSW (z: v, {) < 13 nR
STD
(8)
This result can be further stabilized by taking multiple samples. Thus, we define the Fourier Sliced Wasserstein (FSW) embedding E!ts :PEN (IR") + R by
E (1) = (EES (1) gg()))
FSW
Mi ((() ( ,((m)) g(mm)):
(9)
where (u(*), g(0)) T k=1 are drawn i.i.d. from Uniform (isfoomm(""11 (1 x Dz. Consequently, we have: Corollary 3.3. Under the assumptions of Theorem 3.2,
E E
FtSW (40) EtSw (i)]l SW (M,
(10)
R
STD
1 m
FFS (14) FSW (#)) $ 13 mm"
(11)
Notably, the error bound in (11) is independent of the number of points N and the dimension d. Thus, the estimation error does not suffer from the curse of dimensionality. By taking sufficiently high embedding dimension, (and possibly infinite) support cardinality, while ensuring a bounded standard estimation error, as long the distributions are supported within fixed ball of radius R.
embed distributions of arbitrarily high dimension with arbitrary
4 THEORETICAL RESULTS
In the previous section, we showed that our embedding approximately preserves the sliced Wasser- stein distance probabilistic sense, with diminishing estimation error the embedding dimension grows. Here, we show that with a finite dimension, our embedding is guaranteed to be injective and bi-Lipschitz, as outlined in the Main Results summary in Section
First, we show that with a sufficiently high dimension m, our embedding is injective.
Theorem 4.1. (Proof in Appendix B.3) Let FFESW PPN (R.) - Rm be in (9), with (0(0),((R)) ((0))m k= drawn i.i.d. from Uniform (Sd-1) x De. Then:
1. 1fm 2 2Nd + 1, then with probability 1, FFSW is injective on SCN (R4).
2. mf 2 2Nd 2N then with probability EESW is injective on P&N (R.).
These bounds are optimal essentially up to a multiplicative factor of 2, since an embedding dimen- sion m Nd precludes injectivity for any continuous embedding (Amir et al., 2023, Theorem C.3).
Next, we show that in the case of SeN (Rd). ). the injectivity of FSW implies that it is, in fact, bi- Lipschitz, Our proof relies on the fact that EtSw is piecewise linear and homogeneous in the input points, sense shall now define. By slight abuse of notation, we denote the distribution parametrized by the points X (1x(1) .-... (N)) and weights (W1,.-.;uy WN) (X, un)

---

Published as a conference paper at ICLR 2025
Definition. Let E D - R with D P&N (R.) or SEN (R.). We say that E is positively homogeneous if for any 2 O and any distribution (x, w) E 'D E(aXX, w) a aF w).
The next theorem shows that any injective embedding of PsNN (R*) that is positively homogeneous and piecewise linear* is bi-Lipschitz when restricted to distributions with fixed weights.
Theorem 4.2. (Proof in Appendix B.3) Let E PEN (R2) + mm be injective and positively homogeneous. Let w i E AN be fixed, and suppose that the functions E(X, w) and E(X, ()) are piecewise linear in X. Then there exist C, c >> 0 Such that for all X, X E IRdx and E [o],
We (X, w) X, ))) - E(X, w)) E x) - xx
S. W/ ( (X,, w)) (x, (x) (12)
The restriction to fixed weights can easily be relaxed to weights from any finite set. Based on this observation, we now show that EESw is bi-Lipschitz on multisets.
Corollary 4.3. Let EFSW be as in (9) with m 2 2Nd + 1. bi-Lipschitz on SgN (TR)).
Then with probability I, ESS is
m
Proof. Any μ E SCN (R can be represented by a parameter of the form (X, w()),, where
w (n)
Is S N.
(13)
For n r E [N], let Car, Car be the Lipschitz constants C' of (12) for FES with the weight vectors w w t (t), Then, it follows directly that EFSW is bi-Lipschitz on SgN (R²) with the constants e = min Car and C = maxn,re[N] Car < CC.
Next, explore whether it is possible to further improve by finding bi-Lipschitz embedding for the whole P (R*). For the broader class of distributions UNEN Psx (I=). Naor & Schecht- man (2007) proved that no bi-Lipschitz embedding exists into ([0, 1l), and thus not into any finiteddimensional space. One may ask whether this impossibility can be circumvented by bounding the number of support points and confining them to fixed compact domain, namely, by restricting to N (52) for compact
Rd. The following theorem shows that even this restricted collection still cannot be embedded in bi-Lipschitz manner into any finiteddimensional Euclidean space.
Theorem 4.4. (Proof in Appendix B.3) Let E PPN (S.) - I. where N 2 2 and Ω Rd has nonempty interior. Then for all p E [1, 20], E is not bi bi-Lipschitz on P&N (9) with respect to Wp.
5 NUMERICAL EXPERIMENTS
In this section, we demonstrate how the theoretical strengths of our method manifest in practice. Specifically, we show that our method produces embeddings with superior distance preservation and improves performance in practical learning tasks.
Comparison with PSWE This experiment compares our method with PSWE (Naderializadeh al., 2021) an embedding method developed with the same purpose of preserving the sliced Wasserstein distance. As mentioned in Section 2.2, this method embeds continuous distributions into an infiniteddimensionall Hilbert space, and does so isometrically with respect to the sliced Wasser- stein distance. Finite multisets and distributions are treated by this method as discretizations of continuous distributions, which causes different finite distributions to be identified with the same continuous distribution, consequently leading to a loss of injectivity.
To demonstrate this, we constructed a pair of multisets X1, X2 C 1R, of sizes 5 and 200, respec- tively, that PSWE identifies the same continuous distribution, and thus fails to distinguish; see Ap- pendix C. for details. Table reports the normalized Euclidean distances y [[)) -(x,))
where E SeN (IR3) + I is either FSW or PSWE, for various embedding dimensions m. As shown in the table, PSWE fails to distinguish the two multisets for all tested values of m, whereas FSW succeeds even when m = 1. Moreover, as m increases, the distances produced by FSW converge to the sliced Wasserstein distance SW(X1, X2), as promised by Corollary 3.3.
A definition of piccewise linearity appears in Definition B.12 Appendix B.3.
8

---

Published as a conference paper at ICLR 2025
Table 1: Embedding distance vs. embedding dimension
Embedding dimension
Embedding distances between pair of multisets X1, X2 R distinguished by FSW but not by PSWE. Here [X] = 5, [X]
Method
10
100
1000
10000
FSW
0.125
0.0326
0.0639
0.0733
0.0737
PSWE
0.0
5,48e-8
1.47e-7
1.3e-7
1.44e-7
200, and 1,, =
0.0734.
Learning to approximate the Wasserstein distance One possible approach to overcome the high computation time of the Wasserstein distance for d > 1 is try to estimate it using a neural network, trained on pairs of point-clouds for which the distance is known. This approach was used in pre- vious works (Chen & Wang. 2023; Kawano et al., 2020), which proposed architectures designed to approximate functions F: SeN (R+) (cx((R + R. such the Wasserstein distance function. These methods handle multisets using the traditional approach of sum- or average-poolingg
Since embedding is bi- bi-Lipschitz with respect to the Wasserstein distance, it is a priori likely to be more effective building block for architectures designed to learn it.
For this task. we used the following architecture: First, FSW embedding Eg PPNN(R") IRT is applied to each of the two input distributions μ., p. Then, a second FSW embedding Eg Se2 (RR - R is applied to the multiset (Ei Ea (z)). The output of Eg is then fed to MLP + R + R see Appendix C.2 for dimensions and technical details. Our full architecture described by the formula
F(u, (t) == ((2(((E (())))
This formulation ensures that F is symmetric with respect swapping μ and j. In addition, used leaky-ReLU activations and no biases in 4 which renders F scale-equivariant by design, i.c.
F((xX, un), (a:x, ii) (FX,((,,)) (X; u)) Ya
as is the Wasserstein distance function, which F is designed to approximate.
The experimental setting was replicated from (Chen & Wang. 2023). The objective is to approxi- mate the 1-Wasserstein distance Wi. The following evaluation datasets, kindly provided to us by the authors, were used: Three synthetic datasets no sy soo.sy-sphhe--eee re-3, noisy-sphere- and uni form, with random point-clouds in R3, IR. and IR2;; two real datasets ModelNet-smi 11 and Mode with 3D point-clouds sampled from ModelNet40 objects (Wu et al., 2015); and the gene-expression dataset RNAseq (Yao al., 2021), with multisets in 2000
We compared our architecture to the following methods: (a) VSDeepSetss a DeepSets-like architecture trained to compute Wi-preserving embeddings, and MProodutNet. which further processes the sum of the two embeddings through an MLP (Chen & Wang, 2023); (b) a Siamese autoencoder called Wasserstein Point-Cloud Embedding network (WPCE) (Kawano et al., 2020); and (c) the Sinkhorn approximation algorithm (Cuturi, 2013). We also evaluated the PSWE embedding by employing in our architecture instead of Ej-
Table 2: W-sssersteein approximation: Relative error
Dataset
d
size
Ours
PSWE
NanndurNa
WPCE
NrDeepSten
Sinkhorn
3
100-299
1.4 %
2.2 %
4.6 %
34.1 %
36.2 %
18.7 %
noisy-sphere-6
100-299
%
1.4%
1.5%
26.9 %
29.1
13.7
uniform
256
2.4 %
2.1 %
9.7% %
12.0 %
12.3%
7.3 %
Momallettsmall
20-199
2.9
5.7%
8.4
7.7
10.5
10.1
MedelNet-large
2047
2.6 %
2.4
14.0% %
15.9 %
16.6%
14.8 %
RNAseq
2000
20-199
1.1 %
1.2
1.2%
47.7 %
48.2
4.0
Mean relative emror in approximating the 1-Wasserstein distance between point sets.
As seen in Table 2, our architecture achieves the best accuracy on most evaluation datasets. Training times are in Table 3. Further details appear in Appendix C.2.

---

Published as a conference paper at ICLR 2025
Table 3: 1 L-Wasserstein approximation: Training time
Dataset
Ours
PSWE
PrrouceNet
WPCE
NeDeppSes
noisy-sphere-3
2.2 min
33 min
6min
h 46 min
9 min
sphere-
min
1h
12 min
min
1 38 min
uniform
3 min
51 min
3h 36 min
1 27 min
ModalNet-small
3 min
48 min
min
1 23 min
12 min
ModelNet-large
14.2 min min
min
min
3 5 min
40 min
RNAseq
50 min
15 min
14h 26 min
3 I min
Training times for the different architectures.
Robustness to parameter reduction Next, we evaluate our em- bedding on the Mode 40 object classification benchmark (Wu et al., 2015). This dataset consists of 3D point clouds representing objects coming from 40 different classes.
100%
80%
60%
To demonstrate the improved representational power of our method, we used simple architecture the form F(u) 4b (p)), where E SCN (Rd) R" is the FSW embedding, 4 : Rm - RC is an MIP and C = 40 is the number of classes. We compared this architecture with PointNet (Qi et al., 2017a), which applies an MI.P elementwise, followed by max pooling and another MLP. We eval- uated different sizes both architectures to assess their robustness to parameter reduction.
40%
20%
Point
0%
104
10
10"
Number of parameters
Figure 2: ModeINet40 classifica- tion accuracy
With 300 parameters or more, the two methods perform comparably: our model achieves 85.7 % 86.8 % accuracy and PointNet achieves 84.4 85.6 %. However, model is considerably more robust with fewer parameters, as shown in Figure 2. For example, with k parameters, our model achieves 83.5 % accuracy, whereas PointNet achieves only 54.1 % with 35 k parameters.
We note that simple methods such as PointNet and ours do not achieve the best results on Mod- elNet40. More complex methods achieve superior results by applying PointNet-based embeddings to local neighborhoods of each point. For example, PointNet++ (Qi et al., 2017b) achieved 90.7 % accuracy, and (Mohammadi al., 2021) achieved impressive 95.4 %. Based on our preliminary results, we believe that combining our embedding with such methods, by employing it in their mul- tiset aggregation steps, may enhance robustness to parameter reduction. This property can be crucial in practical applications.
6 CONCLUSION
In this paper, we introduced the first injective embedding for multisets and measures that is provably bi-Lipschitz on multisets. Our embedding is computationally efficient, and both theoretical and empirical results indicate that it better preserves the original geometry of the data, consequently leading to improved performance practical learning tasks.
Our embedding has already shown promising results in subsequent work (Sverdlov et al., 2024), where was used to construct the first graph neural network with bi-Lipschitz Weisfeiler-LLemnn (WL) separation power.
Future research directions include further exploration of the FSW embedding as a building block for learning tasks involving non-Euclidean data, as well extensions to broader notions of distance, such as partial or unbalanced optimal transport.
Reproducibility Statement All experiments in this paper are fully reproducible. The code for training and evaluation, along with the datasets and actual numerical results presented in this paper,
'These results were obtained using standard PyTorch implementation of PointNet (Xia, 2019). The origi- nal result of PointNet's TensorFlow implementation is 89.2 % accuracy.
10

---

Published as a conference paper at ICLR 2025
available at the URL https: http::///gtthhi did not use fixed random seeds for our experiments, the results are consistent across multiple runs. Appendix C provides further technical details regarding experiments.
While
The authors are partially funded by ISF grant 272/23. We thank Samantha Chen for her assistance in reproducing the experiments from (Chen & Wang, 2023).
Acknowledgements
REFERENCES
Miik Aittala and Fredo Durand. Burst image deblurring using permutation invariant convolutional neural networks. In The European Conference Computer Vision (ECCV), September 2018.
Jason Altschuler, Jonathan Niles- Weed, and Philippe Rigollet. Near-linear time approximation al- gorithms for optimal transport via Sinkhorn iteration. Advances neural information processing systems, 30, 2017.
Tal Amir, Steven Gortler, Ilai Avni, Ravina Ravina, and Nadav Dym. Neural injective functions for multisets, measures and graphs via a finite witness theorem. Advances in Neural Information Processing Systems, 36:42516-4255 2023.
Radu Balan and Efstratios Tsoukanis. G-invariant representations using coorbits: Bi-Lipschitz prop- erties, 2023a.
Radu Balan and Efstratos Tsoukanis. Relationships between the phase retrieval problem and per- mutation invariant embeddings. In 2023 International Conference Sampling Theory and Ap- plications (SampTAA), pp. 1-6. IEEE, 2023b.
Radu Balan, Naveed Haghani, and Maneesh Singh. Permutation invariant representations with ap- plications to graph deep learning. arXiv preprint M2Xiv:22033075460 2022.
Erhan Bayraktar and Gaoyue Guo. Strong equivalence between metrics of Wasserstein type, 2021.
Mary L Boas. Mathematical methods in the physical sciences. John Wiley & Sons, 2006.
Nicolas Bonneel, Julien Rabin, Gabriel Peyré, and Hanspeter Pfister. Sliced and Radon Wasserstein barycenters of measures. Journal of Mathematical Imaging and Vision, 51:22-45, 2015.
Jameson Cahill, Joseph Iverson, Dustin G Mixon, and Daniel Packer. Group-invariant max filtering, arXiv preprint arXiv:2205.140399 2022.
Jameson Cahill, Joseph W. Iverson, and Dustin G. Mixon. Towards bilipschitz invariant theory, 2024.
Samantha Chen and Yusu Wang. Neural approximation of Wasserstein distance via a universal archi- tecture for symmetric and factorwise group invariant functions. Advances in Neural Information Processing Systems, 36:9506-999 9517, 2023.
Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. In C.J. Burges, L. Bottou, M. Welling, Z. Ghahramani. and K.Q. Weinberger (eds.), Ad- vances in Neural Information Processing Systems, volume 26. Curran Associates, Inc., 2013. URL https::/proocedddngg hww.rddinnnnnnnennseeeeeee.e.ree fll.aar..osshaaaaaaaaasafffeeee2
Pavel Dvurechensky, Alexander Gasnikov, and Alexey Kroshnin. Computational optimal transport:
Complexity by accelerated gradient descent is better than by Sinkhorn's 's algorithm. In Interna-
tional conference on machine learning, pp. 1367-1376. PMLR, 2018.
Naday Dym and Steven Gortler. Low-dimensional invariant embeddings for universal geometric learning. Foundations of Computational Mathematics, pp. 1-41, 2024.

---

Published as a conference paper at ICLR 2025
Rémi Flamary, Nicolas Courty, Alexandre Gramfort, Mokhtar Z. Alaya, Aurélie Boisbunon, Stanis- las Chambon, Lactitia Chapel, Adrien Corenflos, Kilian Fatras, Nemo Fournier, Léo Gautheron, Nathalic T.H. Gayraud, Hicham Janati, Alain Rakotomamonjy, levgen Redko, Antoine Rolet, Antony Schutz, Vivien Seguy, Danica Sutherland, Romain Tavenard, Alexander Tong, and Titouan Vayer. Pot: Python optimal transport. Journal of Machine Learning Research, 22(78): 1-8, 2021. URL http: / jmlr org/ httpp/pppprs/////2 20 200
Justin Gilmer, Samuel S. Schoenholz, Patrick F. Riley, Oriol Vinyals, and George E. Dahl. Neu- ral message passing for quantum chemistry. In Doina Precup and Yee Whye Teh (eds.), Pro- ceedings of the 34th International Conference on Machine Learning, volume 70 of Proceed- ings of Machine Learning Research, pp. 1263-1272. PMLR, 06-11 Aug 2017. URL https: //proceedings.. .ml.g..ml...........
.mlr. www.pprs/svvv. mmeeela.aa... html.
Branko Grünbaum. Convex Polytopes, volume 221. Springer Science & Business Media, 2003.
Doron Haviv, Russell Zhang Kunes, Thomas Dougherty, Cassandra Burdziak, Tal Nawy, Anna Gilbert, and Dana Pr'Er. Wasserstein wormhole: Scalable optimal transport distance with trans- formers. ArXiv, 2024.
Piotr Indyk and Nitin Thaper. Fast image retrieval via embeddings. ICCV '03: Proceedings of the 3rd International Workshop on Statistical and Computational Theories of Vision, 2003.
Frank Jones. Lebexgue integration on Euclidean space. Jones & Bartlett Learning, 2001.
Keisuke Kawano, Satoshi Koide, and Takuro Kutsuna. Learning Wasserstein isometric embedding for point clouds. In 2020 International Conference 3D Vision (3DV), pp. 473-482. IEEE, 2020.
Thomas N Kipf and Max Welling. Semi-supervised classification with graph convolutional net- works. In International Conference on Learning Representations, 2016.
Soheil Kolouri, Se Rim Park, and Gustavo K Rohde. The Radon cumulative distribution transform and its application to image classification. IEEE transactions on image processing. 25(2):920- 934, 2015.
Juho Lee, Yoonho Lee, Jungtaek Kim, Adam Kosiorek, Seungjin Choi, and Yee Whye Teh. Set transformer: framework for attention-hased permutation-invarian neural networks. In Interna- tional conference machine learning, pp. 3744-3753. PMLR, 2019.
Yuzhe Lu, Xinran Liu, Andrea Soltoggio, and Soheil Kolouri. Slosh: Set locality sensitive hash- ing via Waseeddssssrrttin embeddings. In Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision, pp. 2566-2576, 2024.
Haggai Maron, Heli Ben-Hamu, Hadar Serviansky, and Yaron Lipman. Provably powerful graph networks. Advances in neural information processing systems, 32, 2019.
Haggai Maron, Or Litany, Gal Chechik, and Ethan Fetaya. On learning sets of symmetric elements. In Hal Daumé III and Aarti Singh (eds.), Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pp. 6734-6744. PMLR, 13-18 Jul 2020. URL https: http://cccceccn.e.c/ee/ mlr www.pppesseee.......... sesess/vl99/ m90.m//m0.0......mm html.
Seyed Saber Mohammadi, Yiming Wang, and Alessio Del Bue. Pointview-gen: 3d shape classifica- tion with multi-view point clouds. 202/ IEEE International Conference on Image Processing (ICIP), 3 3103-3107. IEEE, 2021.
Navid Naderializadeh, Joseph F Comer, Reed Andrews, Heiko Hoffmman,, and Soheil Kolouri. Pool- ing by sliced-W Wasserstein embedding. Advances in Neural Information Processing Systems, 34: 3389-3400, 2021.
Assaf Naor and Gideon Schechtman. Planar carthmover is not in L1. SIAM Journal on Computing, 37(3):804 826, 2007.
James Orlin. A faster strongly polynomial minimum cost flow algorithm. In Proceedings of the Twenlieth annual ACM symposium on Theory of Computing, pp. 377-387, 1988.
12

---

Published as a conference paper at ICLR 2025
Sergey Pozdnyakov and Michele Ceriotti. Smooth, exact rotational symmetriza- tion for deep learning on point clouds.
A. Oh, Naumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine (eds.), Advances in Neural Information Pro- cessing Systems, volume 36, pp. 79469-79501. Curran Associates, Inc., 2023. URL https://proceediigg www.nngg...nuurpeee neurips ww/paper__pp..e..ep....eeee... ffrffilles /ta.felss//aaer22233//ii.t Paper- Conf-reneeeeeeccceeeee
Charles R. Qi, Hao Su, Kaichun Mo, and Leonidas J. Guibas. PointNet: Deep learning on point sets for classification and segmentation. Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR), July 2017a.
Charles Ruizhongtai Qi, Li Yi, Hao Su, and Leonidas Guibas. PointNet++: Deep hierarchical fea- ture learning on point sets in a metric space. Advances in neural information processing systems, 30, 2017b.
Yonatan Sverdlov, Yair Davidson, Nadav Dym, and Tal Amir. FSW-GNN: A bi-Lipschitz WL- equivalent graph neural network, 2024. URL
2410. 011.1991..11.. 09118.
Puoya Tabaghi and Yusu Wang. Universal representation of permutatioon-nnvariant functions vec- tors and tensors. In Claire Vernade and Daniel Hsu (eds.), Proceedings of The 35th International Conference on Algorithmic Learning Theory, volume 237 of Proceedings of Machine Learning Research, pp. 1134-1187. PMLR, 25-28 Feb 2024. URL https: htt.httss///rrrreeedngge pttp///rreeeen...cooooo/oo .ma.eeeddings............ mmlr. press/v2333 http:///////aaaaaa/eeeeee httpt/aaaaaaaaaaaaaaaa html.
Petar Veličković, Guillem Cucurull, Arantxa Casanova, Adriana Romero, Pietro Liò, and Yoshua Bengio. Graph attention networks. In International Conference on Learning Representations, 2018.
Cédric Villani al. Optimal transport: old and new, volume 338. Springer, 2008.
Edward Wagstaff, Fabian Fuchs, Martin Engelcke, Ingmar Posner, and Michael A Osborne. On the limitations of representing functions on sets. In International Conference on Machine Learning, pp. 6487-6494. PMLR, 2019.
Edward Wagstaff, Fabian B Fuchs, Martin Engelcke, Michael A Osborne, and Ingmar Posner. Uni- versal approximation of functions on sets. Journal of Machine Learning Research, 23(151):1-56,
2022.
Larry Wasserman. All of statistics: A concise course in statistical inference, volume 26. Springer, 2004.
Zhirong Wu, Shuran Song, Aditya Khosla, Fisher Yu, Linguang Zhang, Xiaoou Tang, and Jianxiong Xiao. 3d shapenets: A deep representation for volumetric shapes. Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 1912-1920, 2015.
Fei Xia. PointNet.pytorch. https://github. http://ttttuuu..coooi////2/2// http://hhhthc..cmom/ma/2//tttttt pytorch, 2019.
Keyulu Xu, Weihua Hu, Jure Leskovec, and Stefanie Jegelka. How powerful graph neural networks? In International Conference on Learning Representations, 2018.
Zizhen Yao. Cindy TJ Van Velthoven, Thuc Nghi Nguyen, Jeff Goldy, Adriana E Sedeno-Cortes, Fahimeh Baftizadeh, Darren Bertagnolli, Tamara Casper, Megan Chiang, Kirsten Crichton, al. A taxonomy of transcriptomic cell types across the isocortex and hippocampal formation. Cell, 184(12):3222-324 2021.
Manzil Zaheer, Satwik Kottur, Siamak Ravanbhakhsh, Barnabás Póczos, Ruslan Salakhutdinov, and Alexander Smola. Deep sets. In Proceedings of the 31st International Conference Neural Information Processing Systems, NIPS' 17, pp. 3394-3404 Red Hook, NY, USA, 2017. Curran Associates Inc. ISBN 978151086049
13

---

Published as a conference paper at ICLR 2025
Muhan Zhang, Zhicheng Cui, Marion Neumann, and Yixin Chen. An end-to-end deep learning architecture for graph classification. In Sheila A. Mcllraith and Kilian Q. Weinberger (eds.), Pro- ceedings of the Thirty-Second AAAI Conference innovative Applications of Artificial Intelligence (IAAI-18) 18), and the 8th AAAI Symposium on Ed- ucational Advances in Artificial Intelligence (EAAI-18), New Orleans, Louisiana, USA, Febru- URL
Artificial Intelligence, (AAAI-18), the 30th
ary 2-7, 2018, pp. 1438-4444 AAAI Press, 2018. doi: 10.160/AAALV332111220
https: www.i.orgg/......
aaai w32.1.
Yan Zhang, Jonathon Hare, and Adam Prügel-Rennett. FSPool: Learning set representations with featurewise sort pooling, arXiv preprint (eXXx::906.02999, 2019.
14

---

Published as a conference paper at ICLR 2025
The appendices organized as follows: Appendix A provides additional details the FSW em- bedding beyond those covered in the main text. Appendix provides proofs for our theoretical results. Appendix C presents technical details of our numerical experiments.
A FURTHER DETAILS ON THE FSW EMBEDDING
In this appendix, we discuss additional aspects of the FSW embedding. In Appendix A.1, we extend Em from probability distributions to measures of arbitrary total mass while retaining injectivity. In Appendix A.2, we present explicit formulas for practical computation and analyze the computational complexity of the embedding.
A.1 EXTENSION TO MEASURES WITH ARBITRARY TOTAL MASS
Here, extend the definition of the FSW embedding measures that not necessarily probabil- ity measures-that is, measures whose weights do not sum to 1-while maintaining injectivity. This extension is crucial for distinguishing multisets of different sizes when their element proportions are identical (as in footnote 3, Page 4), as well as for ensuring continuity when handling the empty multiset as input.
This is particularly important when using the embedding as message aggregation function in graph neural networks, since (1) in graph-learning tasks, vertex neighborhoods of different sizes but with the same neighbor proportions are typically treated as distinct, and (2) isolated vertices must aggre- gate messages over the empty multiset. In common graph architectures that require weight normal- ization, this issue is usually addressed by adding self-loops to ensure that vertex has empty neighborhood. The solution presented here eliminates the need for such modifications to the graph structure.
Denote by McN(2) the collection of all measures with at most N support points over a measurable domain f2 C Rd. Any measure μu. E M (S2) can be represented by points x(1), ,=(,,...,,a" (N) E Ω
and corresponding weights wi....; WN 2 weighted sum of Dirac's 's delta functions
&
(14)
Definition. Let y be as in (14). The total mass of μ is defined as
((2)
Wi-
A first step towards extending FFS from P&N to M2)((2) while maintaining injectivity is to simply add an extra output coordinate that encodes the total mass of the input measure. Define frsw :McN(() \ 0 + Inm for an input μ as in (14) by
ErSw (1) = [p(S2), Etsw (P)].
(15)
where j is the measure μ normalized have a total mass of 1:
N
Wi
6 Sz(1,
(16)
1= Wy
easy to verify that with the above definition, by Theorem 4.1, if m 2 2Nd + 2N then Ærsw
" is injective on McN (Rd) \ 0, i.e., the collection of all measures in Mgx excluding the zero measure μ 0, for which j is not 5 Moreover, it further follows from the theorem that m 2Nd 2 suffices to ensure that WrSw distinguishes between distinct pairs of input multisets of size < N, even if their element proportions are identical and they differ only in their cardinality.
SStrittyy speaking, it is injective for almost any choice of the embedding parameters (z()), ((()))(()) in the statement of Theorem 4.1. We ignore this here for brevity.
as
km
15

---

Published as a conference paper at ICLR 2025
One limitation of the definition in (15) is that Emw is not defined the zero measure 0 and, moreover, has jump discontinuity there. This issue can be remedied by padding imput measures whose total mass is below chosen threshold, assigning the complementary mass to the zero vector in IRd. The result is then normalized to a probability distribution, as before.
Namely, choose threshold p 0 and define the regularized measure pp for MgN (R.) by
(( Z P.
Pp
(()
80
(1) < p.
(17)
where dg is Dirac's delta function at Rd.. We then redefine EESW to use Up rather than ja:
E (p) [m(2), E (1p))
(18)
With the definition in (18), Etsw well-defined on the whole of Mgg (R*). Moreover, it is tinuous and has finite directional derivatives everywhere, including the zero measure. As discussed above, if m 2 2Nd + 2N, then "m suffices.
FESW is injective, and for injectivity on multisets, m 2 2Nd +
Note that does not have to be small. In fact, very small values of p lead to large gradients of FFSW near the zero measure. In our experiments with a message-passing neural network based on the FSW embedding (Sverdlov et al., 2024), we found that values the range 0.1 S p S 1 work well in practice.
Maintaining homogeneity In some settings, is desirable to have an embedding that is positively homogeneous with respect the input points? while still encoding the input multiset size. This property is particularly useful in architectures designed to approximate multiset functions that are known to be homogeneous, yet required distinguish multisets of different sizes and accommodate the empty multiset.
Note that this necessitates slight compromise of the injectivity, as any such embedding inevitably maps all measures supported on the zero vector 0 E Rd to the zero vector 0 E R", regardless of their total mass.
To achieve this homogeneity, straightforward solution is to use slight modification of (18) and multiply the first coordinate p(22) by the norm the internal embedding FFS (ptp):
B.s (4) || 1 (up))) ((2)), BESW (H1)
(19)
With the formulation (19), the positive homogeneity of EFSW, implies that EFSW is also positively homogeneous. Moreover, it can be verified that Ersw with this formulation still distinguishes be- tween any pair of distinct input measures, except for the case that both measures are supported at 0 E ddd.
A.2 PRACTICAL COMPUTATION
Here we present formulas that facilitate the practical computation of EESw and then analyze its computational complexity.
We start by introducing notation that shall be used to express quantile functions of distributions in PsN(R)
Definition A.I. For vector æ (I (II...., xN)) N the order statistics .(1).... J(N) are the coordinates of æ sorted in increasing order: T(() ... S M(N). We define the sorting permutation
a (a (z:), ...) o xxx) E Sx
to be permutation that satisfies (x) = vis) for all i [N], with ties broken arbitrarily.
See definition Section
16

---

Published as a conference paper at ICLR 2025
We now show how (t) can be expressed explicitly in terms of the order statistics of the support of μ. Let
p wist PSN
and denote (#1, (E1...., IN), w (W1,... .... wx). Then for all t E [0, 1), it can be shown that
(t) (o()),(((
(20)
where kimin (a, w. 1) is defined for (0) 101....,0N: ON) SN by
kmin(a, w, t) = min { E [NT] wo, +... Wo, t}. It can be seen in (20) and (21) that (t) is monotone increasing with respect to t. Moreover, 2₂ (0) ess min (u) and lim Qx(() max (u),
(21)
with ess min (μ) and ess max (μ) denoting the essential minimum and maximum of the distribution (a. We thus augment the definition Q₂ to [0, 1] by setting Q₂, = ess max (n).
[0. 1]. is continuous at both end points, and since it is monotone increasing, it only has jump discontinuities.
Note. With the above definition, Qu is right-
Using the identity (20), E(u; v. can be expressed as
Σ (x(
(x)) Qv cos (2xtt)dt
N
(x)xxx
-2(1 + g)
(x"x)
x) (1) cos (2rgt)dt
(22)
N=
(xxxx
N
sin (2Tet) Eiz (o,("(')) It=l (xX()
-2 2
(16)
under the notion Ei=1 W Wo = 0. Rearranging terms gives us the alternative formula
1 1 Σ
[00xx (1) (o= (a)))
E(u; e,
sin 2 Wo x)
(23)
2
with the definition of (v x) extended to A: N 11
(x'X) (N+1) (N+))
To prevent numerical instability for values of & zero, the explicit division by 2 can be avoided by replacing the sine function in (23) with the normalized sinc function
(Tx)
I + 0,
sinc(r) =
0,
which is implemented in a numerically stable way in most standard numerical libraries. This yields the alternative formulas to (22) and (23):
Eter (o;((Tx) Σ (v x)(xyl sinc(2Et))
E(p; v, 2(1 + 4
;- D W (x)(
(24)
and
(():,,,)=((11
Σ Amm
Σ wo
sinc 2 Σ wo [10xx) (k) - (o X
(6+1))
(25)
Using the above formulas, the following theorem specifies the complexity of computing EESW (Is
17

---

Published as a conference paper at ICLR 2025
Theorem A.2. The function EFSW (p: (((()))((()))) O'mNd mN log N) time.
Th can be computed in )=1) PSN (R*)
Proof. Computing EFSW (1: (v((k) (())))=)))) consists of the computation of EFSW (n;;
independently for k ,,..., m. In (23) it can be seen that the most expensive operations in each such computation are: (1) the computation of inner products in IR and (2) finding the sorting permutation of these inner products. Therefore, each such computation takes O + N log N) time.
18

---

Published as a conference paper at ICLR 2025
B PROOFS
This appendix contains proofs for the theoretical results presented in the main text. Appendix B.2 provides proofs for our probabilistic results Theorem 3.2 and Corollary 3.3, which rely on prop- erties of the cosine transform discussed in Appendix B.1. Readers interested only in injectivity and bi-Lipschitzness may skip directly to Appendix B.3, which contains proof's for the results from Section 4.
THE COSINE TRANSFORM
The cosine transform takes major role in proofs. We now define it and present some of its properties. The results in this subsection appear in standard textbooks such as (Jones, 2001; Boas, 2006). We include them here for completeness.
In the following discussion, _ denotes the space (P(R), defined by
(P(R) = {f R + R |
is Lebesgue measurable and [l < o0),
with
[sn 15(0) di) 1/p css SUPIER f())
p (0)
p p= 00.
Definition B.1. Let f E L such that f(1) 0 for all / < 0. Define the cosine transform of f is
f(E) = /0- f(t) cos (2rtt)dt
(26)
for & 2 0.
Note that if f E L,, then
|101 2211fll
(27)
since
[(() $2 $2/1 [f()) |cos (2T $2 ( [f(t) fll
(28)
Thus, if f L,,, then f L The following lemma proves a better bound as E + 00 if f is monotonous, and shows that the cosine transform preserves the L²-norm.
Lemma B.2 (Properties of the cosine transform). Let f L! such that f(t) for all 0. Then:
1. L' then
[ (f(t)) - (P(()) dt.
(29)
2. Suppose that f L and that f is monotonous on an interval I =0,T) and vanishes almost everywhere outside of I. Then for any & > 0,
3 x lf
f(E) <
(30)
Proof. We start from part 1. Throughout the proof, we consider the natural extension of f({) to negative values & according to (26), namely
-(-E)= = f(E).
Let fe(t) be the component of f.
fa(() +0((0+++++))))) 25(()).
=)(())
19

---

Published as a conference paper at ICLR 2025
Then
To(5)
Se(1)e 2xiEt dl (
fe()) (-2xE()dt
{(((t) + f(-t)) cos (-2
((t)+f(-t)) cos (-2 + [ (f(t)+ cos ( (-2r
f(-t) cos (-2 /- f(t) cos (-2xtt)dt
f(r) cos (2mtr) + 2 / f(t cos (2xtt)dt
f(t) cos ( =
with (a) holding since the Fourier transform of real even function is real. Hence,
f(E) = =(F(),
& R.
(31)
Now, J, (f(t) de
2 (f(t) 2 dg
1
2/2 $1| 2
(2) 2
Pal 2| @ 11.1
(f()) L. dt
where (a) is by (31), (b) holds by the Plancherel theorem, and (c) holds since
(t)( dt L. (b(f(t) f(- f(-t))) dt [((((0)) + -100((-)) + ((((-))))
Il 122
[1OO) + +(((-1))]a) di
((()) + +1 $1 (s(-t))²
| (S()) = 1 (((1))2l = 11
and thus part holds. We now prove part 2.
Suppose first that f is differentiable on J. Using integration by parts, we have
f(E) =(E)=2 2 = //' f(t) cos (2mtt)dt
1= E f f(() sin (2TEt) dt.
f(t) sin (2TEt) (inT(Et)))
nt
Let us now bound the terms A and Ag
|4| (T( sin (2mET) f(0) of < IS(T)) |
20

---

Published as a conference paper at ICLR 2025
and
T
A2 =
(t( sin (2TEt)dt
! [5"(1)
((00)-|siin |sin (2TEt)ldt
< / (f) ( !.." ((()dt
[T
II(TT) ((0)) S 2
with (a) holding since f" does not change sign (0, 7) due to the monotonicity of In conclusion, we have
(())
(1 S x
1 A2l) S T
(32)
To remove the differentiability assumption on f. we use the technique of mollification, and replace by sequence of smooth functions that converges in L'; see Chapter 7, Section C.3 of (Jones, 2001).
For the smooth functions to be monotonous, we first define a modified function f: + R
f(0o)
t t
f(1) = f(I) LE (0,7)
(33)
F(T T2
where s(0+) and S(T-) are the one-sided limits of f(t) as l approaches 0,7 T respectively from within the interval T). Note that the monotonicity of f. combined with the fact that f L implies that both limits must be finite.
Now, f coincides with f on I, is monotonous on I, and by construction,
| |f
Let Φ= : R - R for e > be the mollifying function defined (Jones, 2001), page 176. We now list few properties of Φs.
1. de is infinitely differentiable and compactly supported.
2. de is radial, i.e. (e(() = (-(-)..
3. (t) 2 0 for all t, and (t) 0 iff |t| <E
4. Jg de (t)dt=
Let fe R R for e > 0 be defined by
ft) = xi(t) / (t-)
/ f(t (r)dr,
(34)
with Xı denoting the characteristic function of I. From the rightmost part of (34), it is evident that the monotonicity of f implies that fe is monotonous on I.
Also note that
[a(() S xr(1) / j( + ())
$11F] ( (r())r
$
(35)
-||i|| 1 ||
=
21

---

Published as a conference paper at ICLR 2025
Thus,
[ S &TIiflle>
[ s 11
(36)
and hence Je LL From the discussion in (Jones, 2001), Sa satisfies:
1. fe E (I)
2. lime- lIfe - 1 =0
So far we have shown that for any e > 0, fe is in L n L is monotonous and smooth on 7, and vanishes outside of T. Therefore its cosine transform satisfies
(a) 3 S T
(b) 3
[ell L S
x |l
(37)
where (a) by (32) and (b) is by (36). Thus,
I (fe(t) f(t)) cos (2mtt)
f - f(o)
Se 5لاا ||cos (2m
<I fll.: = 0.
In conclusion,
3 |l
(m7) SP(le)
=ot f(()
2
and therefore
[f(e)
S TE
|| 1
B.2 PROBABILISTIC PROPERTIES
In this subsection, we prove our probabilistic claims on EFSW First, we show that FSSW (14 v. E) uniformly bounded for all embedding parameters (v, E) and input distributions μ PeN(M), assuming that Ω Rd is bounded. We then proceed to prove Theorem 3.2.
For the discussion below, we extend the definition of quantile functions Q, : [0. 1) - to functions on the whole interval Q₂ IR - R. using the convention (t,(t) = for 00, 1). and defining (1) to be the limit lim; (t), which is known to exist, since Q₂ is bounded and monotonous.
Let us start by defining the notation
A(U, ū: v, == 1SW (, ESW ("
(38)
We define pseudonorm for distributions in PSN (R'):
IIll,, = Wp(p, 5o). p E [o]
(39)
where dn denotes Dirac's delta function in 0 E Rd, that is, the distribution that assigns a mass to the point 0 Rdd Note that |I- Ilw, is not true norm, since P&x (R is not vector space.
The following claim provides useful bounds on the distance between two distributions and on the quantile function projected distribution, in terms of the above pseudonorm.
Claim B.3. For any p. PAN (IR) and 1d-1.0 1.
SW(,,)) (2) < W(n, ( S 111llwwe + ll
(40)
Qu),((() < ||mllw..
(41)
22

---

Published as a conference paper at ICLR 2025
Proof. The left inequality in (40) is well-known property of the sliced Wasserstein distance; e.g. Eq. (3.2) of (Bayraktar Guo, 2021). The right inequality is easy to verify by considering the transport plans that transport each of the distributions μ., ja to ón and using the triangle for the Wasserstein distance.
Inequality (41) can be proved by showing that the quantile function Qer, only admits values that support points of the projected distribution v² μ. and any such point y IR is inner product y 20 of o with some I E I2 that is support point of μ.
The following lemma shows that E (1:; v, E) is indeed bounded, as mentioned above, and lays the groundwork for proving Theorem 3.2.
Lemma B.4. Let u. J E PSN (RE4) and v E Sd-1 Then
FS (n; [)) 311mlll> VP 0,
(42)
Eg A 2 (1, it; v. {)] (²20"" u. v²=
(43)
STD
W A 2 (1,, it; v, 6) S 3(1lllww + -llill +lllllw))Wo' u, vT)..
(44)
Proof. By the definitions of Etsw and the cosine transform in (6), (26) respectively,
FSS B (u; (1 (1+ ()=
(45)
where QUtU and vanishes outside of q E [0, 1]. by part of Lemma B.2, we have that
n is the cosine transform of the quantile function QUTR. Since Q (t) is monotonous
Quor(()) s 11 1
3 x
(41) 3
x |llllw
(46)
By (28),
(e-r(()) < 22 ull S 21 2 S &2l1allw>>
(47)
with (a) holding since Q a is supported on [0, 1]. Thus, by (46) and (47),
Qurn(E) < min { 2 3 T Illllwww
Combined with (45), the above implies
3 T } Ilellws
3 I (2 = Illllws S 33]lm||ww
(0)
FSW (1: s (1 + 3) min 2.
where (a) can be verified by simple analysis. Thus, (42) holds. Note that since BESW (1 E) is bounded as a function of &, so is A² it; v, E), and therefore both have finite moments of all orders with respect to E.
Now,
Eg A ( (p,
Eg
[ (EESW (1; FFSW (āt; ())
2) " de (Qor)(()) Qora(t))
(1 (L+
(Qurult)-Qur - Qu dg
(Qvza(1) Quta(()) dd
[ (Qor,(t)- - Qur 2 dt
( W2 μ4 u p))
23

---

Published as a conference paper at ICLR 2025
with (a) following from part 1 of Lemma B.2 and the linearity of the cosine transform, and (b) holding by the identity (4). Thus, (43) holds.
To bound the variance of A A² ji; v, E), note that
Vare (H,
E (A (1, pi 0,)))
(u,
(43) Eg $ (1,, ii; v, {)] - (M² u4 x))²
(E RS (c: v, () - ( v 2))² (M, - (0
u, v)
s Eg (LENSW (1o: v. v. + LESS (p: () 2 (1 - W0 10
(2)
Eg (3llmllw + +31l/llw.m.. )² (1, w4 1,
(llullww + -11illww.) )2 (1
w"(0"
91luallw + +llällwa )² ) (0"00" u1 - w (0" (0-M,00'1))
s (llullw.. + +lllllw.. )² w2 (0"11007
and thus (44) holds.
This concludes the proof of Lemma B.4.
Let us now prove Theorem 3.2.
Theorem 3.2. (Statement in Section 3) Let H. E PSN (Rn(BR). with BR C Rdd being the closed
Eg-norm ball of radius R. centered at zero. Let v Uniform Sd ). ).C= Dz. Then
FSW [L (1, - FSW (Ā; v, E) 21 = SW2 (u,
Ev,
(7)
STD 2.
JENSS (1; E (,: {)
21
S 13 R
(8)
Proof. Eq. (7) holds since
Ev.g (1,, i3; v, {)] Evp [Bgje (M,
(43) E. [W2 (0 p)]
(w.
We now prove (8).
Var,,g A (1, R4; u. {)]
(
E. Vare v (1,, ir; n, [)]] + Var, [Egl (4, (,, v, {)]]
E., [arr v A (1, ir: v, {)]] + Var, (v" 11 o² mi]
(44)
E. [(l144|ww + ||Aill )² w (8= u 07 xi)
+ Var. (w2 (o",, v.)"))))))
(IlA|] + +1lälw_)* )²= E. (w2 H, 0 R)] + Vare [M2 (2 (UTT, xi)
5 9(Iullww. + +I1illw.. )² )². sw + Var, (W2 n, o²
9(Ilallw + +Ilillwa)) )² -SW (1, (b) + E. [w (w" u. o Ri)
(40)
< 911m||w> + +11 +11llw-)) )² (Ilullww + =1Alll=)) 2 Ev (1llllww + -1llllwa))
10(1|A|| + l1All
where (a) by (Wasserman, 2004, Theorem 3.27, pg. 55). Now, since I1lllw.. : [l S R,
24

---

Published as a conference paper at ICLR 2025
have
Var (u, = 10(1lll... + +112ll11 )a
I 10 (2R)* = 160 R¹ 13 R²
and therefore (8) holds.
B.3 INJECTIVITY AND BI- BI-LIPSCHITZNESS
Theorem 4.1. (Statement in Section 4) Lel E PPNN + R be as in (9), with
m (w ((()) drawn ii.d from Uniform (Su-1) De. Then:
1. mf 2 2Nd + 1, then with probability 1, Fm is injective on SeN (R4).
2. mf 2 2Nd + 2N 1, then with probability 1, F is injective on PSN (R4).
Proof. Our proof relies on the theory of o-subanalytic functions, introduced in (Amir al., 2023). This theory defines and studies a class of sets called o-subanalylic sets and a corresponding class of functions called o-subanalytic functions. These definitions are technically involved and require an elaborate construction, which we do not include here. However, our proof is self-contained, and we state the relevant properties of these constructs. For a more detailed exposition, we refer the interested reader to (Amir et al., 2023, Appendix A).
The main result of this theory, which plays a central role in our proof, the Finite Wimess Theorem (Amir al., 2023, Theorem A.2). This theorem is powerful tool for proving the injectivity of parametric functions. Essentially, it enables reducing an infinite family of equality constraints to a finite subfamily chosen randomly, while maintaining equivalence with probability 1.
We hegin by stating the full version of the Finite Witness Theorem. The new concepts introduced in the theorem statement will be explained below.
Theorem B.5 (Finite Witness Theorem). Let C IP and W S R be o-subanalytic sets of dimensions D and Da respectively. Let F M W IR be a o-subanalytic function. Define the set
= f2 M F(2: 0) 0, 40 W].
Suppose for all E MI N
dim (0 E W F(2;0) = 0} < Do - 1.
(48)
Then for generic (0(2) ..... O(D-1)) E wD++1
N N f MI
F F(2:0(1) 0, -i= ,D....,,+...
D+13
Moreover, W is an open and connected subset of R2, and F(z;0) is analytic function of 0 for all fixed z E Ml, then condition (48) is not required, as it is automatically satisfied.
The function F(z; 0) in the theorem can be regarded as representing a parametric family of con- straints on z, parametrized by 0: point z E M. belongs to the feasible set N if it satisfies the constraints Fxx; 0) 0 for all parameter vectors
W. The vectors O are called wimess vectors.
The theorem essentially states that, if one draws a sufficiently large number of witness vectors from W at random, then with probability 1, that set of vectors is sufficient to fully determine N.
The number required vectors depends on the dimension of M. The notion dimension used in the theorem is the Hausdorff dimension, which coincides with the standard notion of dimension for vector spaces and smooth manifolds. The condition (48) called the dimension deficiency condition.
Our next step is to show how Theorem 4.1 follows from Theorem B.5, assuming the latter's 's assump- tions are satisfied. We then establish that these assumptions indeed hold. First, we prove part 2 of Theorem 4.1. To this end, must choose the appropriate M. W and F. First, set
M Rdxx
TR Nx
N
25

---

Published as a conference paper at ICLR 2025
where A is the probability simplex in RN Namely, Ms the set of parameters required to describe pairs of distributions μ., J PSN (R*))
M = { x X w. w
x, E RdxN
w, v E AND.
Throughout our proof, we use the convention
A Σ
P
Since dim A -N the dimension of M is
dim M 2Nd + 2N 2.
(49)
Second, set W to be the space of embedding parameters
W = {(V,E)
v E sd-1, E 00)).
Now, define the function
F M x W + R.
F X, w in u, = F. FSW (X, w w; E)- - FFSW (x, ū: w.).
(50)
where, by abuse of notation, ESS (X, w; v, E) denotes ETSW (μ;; U, E) and Ess X w 0; v,E) de- notes EFSW (ji; v. E), with μ., j being the measures defined by (X, w) and (X, w)) respectively.
Essentially, r ((x. X, ui v.)) compares the two input distributions n, by consider- ing their sam-le embeddings E#SW (.) with the parameters v, Specifically, F X, X, w. úi; v. equals zero if and only if both distributions and ğ "appear identical" to EESW with the embedding parameters v.
The full embedding EFSW (-: (u((), (1(())g()))))) )()) uses m pairs of parameters (,(M), > drawn independently random. Recall that part of the theorem, which we wish to prove, assumes that m 2 2Nd + 2N 1.
Combined with (49), the above implies that
2 dim M + 1 111-+1 1
and thus m is a sufficient number of witness vectors as required by Theorem B.5. Therefore, pro- vided that the rest of the assumptions of Theorem B.5 are satisfied, the theorem states that, with
probability on the embedding parameters (v(*), g(())) k=1" the following holds: for any pair of
input distributions pt, j P&N (RRd), FFSWW (-: (z)(A)) g())) (()))m)))
does not distinguish between jr
and j. then
FSW v, = FSW (īi; V,E)
for all parameters (u,E) E W. This, combined with (7), implies that SW(/, (t) = 0, and therefore ja = μ.. Hence, with probability 1, Essw is injective on PsN (Rd).
To prove part of the theorem, namely, injectivity on ScN (R.), we restrict to parameters that describe pairs of multisets in SeN (R.). That is,
M = ( w. (x,,
X. X. E EdN. w, Φ E ww
{
w
(N)
where each (T) for n [N] uniform weight vector corresponding multiset of size n, as defined in (13). Note that with this definition, MI is finite union of affine spaces of dimension 2Nd, and thus
dim M 2Nd.
Recall that part assumes that
m 2 2Nd + dim MI +
and the rest of the proof follows for part
is now left for us to show that F indeed satisfies the remaining assumptions of Theorem B.5, We begin by stating the formal definition of a-subanalytic functions.
26

---

Published as a conference paper at ICLR 2025
Definition B.6. A function f : + B, with A S R B C R", is said be o-subanalytic, if A and B are o-subanalytic sets, and the graph of f. defined by
(f {(x, f(x)) IE T
a o--ubanalytic subset of Rm R
The following proposition presents several properties of o-subanalytic sets and functions that will be used our proof. These results are established in (Amir al., 2023, Appendix A), with each property accompanied by reference to its corresponding statement in that paper.
Proposition B.7 (Properties of o-subanalytic sets and functions).
1. Finite unions, intersections, and Cartesian products of o-subanalytic sets are o-subanalytic (Proposition A.11).
2. Finite sums, products, compositions and concatenations of o-subanalytic functions subanalytic (Proposition A.15).
3. Semialgebraic functions, and in particular, piecewise-linear fimcttion, are o-subanalytic (see Figure 3 therein).
4. Analytic functions defined on open domains are o-subanalytic (Lemma 14).
5. Any o-subanalytic set A can be presented as a countable union
Ci:
where each Ci is C manifold. In particular,
dim A max dim C
(Propositions and A.1
An additional property required for proof is that the preimage of o-subanalytic set under --subanalytic function o-subanalytic. We now prove this.
Proposition B.8. Let f A be o-subanalytic function, and let C C B be a a-subanalytic set. Then the set
(aE
a-subanalytic.
Proof. The set (C) can be presented as
(C) = rGrraph( (4 x C)),
with क. 1 R + R being the orthogonal projection operator
a T( 3)) Thus, follows from Proposition B.7 that (C) o-subanalytic.
Now, argue that M and W o-subanalytic sets. The reason is that both sets semialgebraic, which is easy to verify, and any semialgebraic set is o-subanalytic; see (Amir et al., 2023, Figure 3).
Next, we want show that F' is a-subanalytic function. For this, need show that EtSw (X, w, v, {) is o-subanalytic as a function of (X, N.w, E). To see this, note that by (23), FSw X, w; v. {) is the sum of terms of the form
1+ E sin
(10T)) (k) - (v-x) (
2 > w (o(()Tx)
(51)
2
i=1
27

---

Published as a conference paper at ICLR 2025
with the sum taken over & [M] The product v is semialgebraic, and thus is o- subanalytic. Each term (o x) (u (v*x) (++1) is piecewise linear in v" X and thus d- subanalytic. Each Eiz Wi (x" x) semialgebraic and thus is also o-analytic. The product 2 Ei= "o,("")); composition with the analytic sine function sin (2 "o,(="x)): and again product 2 1+f sin ( "o;("xx)) x)) is also o-subanalytic. Finally, the product (51) and the finite sum of such over E [M] is o-subanalytic. Therefore, F is o-subanalytic, as it the differ- ence of two o-subanalytic functions.
Since F is not generally analytic as a function of (v.E) for fixed ( (x,X, w. (i) we need to ex-
plicitly show that F satisfies the dimension deficiency condition. Let jr, i P&N (R4) be two fixed distributions and suppose that μ + f.. Let A S W be the set
A = (w, E E x (0, oo)
E FSW B (p) = FSW (i; v, {
We need to show that A is dimension deficient; that is, dim < dim W.
It is sufficient to show that A is set of measure zero in W. By measure, we mean that we consider W = 5d x (0, oo) endowed with the product measure, where gd-1 is equipped with the uniform measure and (0, oo) with the Lebesgue measure.
The reason is as follows: A is o-subanalytic since is the preimage of the o-subanalytic set {0} under the o-subanalytic function F F(x, Xx w. w: .). If dim A din W, then by Proposition B.7, XX,X,
A contains some C manifold C such that dim C2 dim W. Any such manifold is nonempty subset of W that is open in the product topology on gd-1 x (0, oc). Since C; is open and nonempty, it has positive measure, and thus does A.
To finish the proof, we now show that A indeed has measure zero in W, namely that almost any (v,E) E W is not in A. First, note that for almost any sd-1, the inner product (v, )) : Rd - IR maps the finite number of support points of μ and j injectively to RR For all such v, the obtained
one-dimensional projected measures will be distinct, that is, v² μ o v² Jı, and therefore, so will be their quantile functions,
QUTN + QUtp.
Once v is also fixed, F (X, X, w. w v.) is analytic as function of &. Thus it must be nonzero almost everywhere. Otherwise, it would be zero at all points &, which would contradict the fact that QuTT n. A Qura and the cosine transform is an isometry.
Thus, we have shown that F satisfies the dimension deficiency condition. This concludes our proof.
Let us now prove Theorem 4.4.
Theorem 4.4. (Statement in Section 4) Let E PgN(()) + RT, where N 2 2 and Ω C IRd has nonempty interior. Then for all p E [1, 20], E not bi- bi-Lihhchitz on Ps,((Q)) with respect to Wp.
Before proving the theorem, note that it implies that most practical embeddings of P (R) likely to fail in SwerrLLpschhtzness since it is reasonable to expect such embeddings to be upper Lipschitz. This is formulated in the following corollary.
Corollary B.9. Under the assumptions of Theorem 4.4, if E (2) - TM is upper-Lipschitz with respect to Wi, then it is not loisc-Lipschitz with respect to any Wp with E [1)
Proof. First, note that it is a well-known fact in optimal transport theory (Villani et al., 2008 Re- mark 6.6) that Wp(n, (u, ū is monotone increasing in p. that is,
Wy(m, & (y,
1 1 q oc.
(52)
Now, if E is upper-Lipschitz w.t.t Wi, then by Theorem 4.4 it is not lower-Lipschitz w.n.t. Wi. Therefore, since Wp(n, ūi) 2 Wi(,, (i) for any p ' 1, E. is thus not lower-Lipschitz w.r.t. Wμ. n
28

---

Published as a conference paper at ICLR 2025
Proof of Theorem 4.4. Our proof consists of three steps. First, in Lemma B.10 below, prove the theorem for the special case that E' is positively homogeneous and Ω is an open ball centered at zero. Then, in Lemma B.11, release the homogeneity assumption by considering homogenized version of F. Finally, we generalize to arbitrary with a nonempty interior in a straightforward manner.
For convenience of notation, before proceeding with the proof, we introduce an alternative scalar multiplication operation for distributions in (S)(
Definition. For μ = Ei= wi0m00 P&N (R*) and a scalar a E IR, we define the distribution oj E (RR) by
ap
w/S az(())))
Note that this definition is not the same as the standard scalar multiplication operation for measures used in Appendix A.1 (e.g., equation (16)), which scales the measure weights rather than the support points. We will use this alternative notation throughout, and only within, this section.
Let us begin with the special case of positively homogeneous E.
Lemma B.10. Let E : P&N (St) + RT with C IR being an open ball centered at zero, N 2 and 2 1. Suppose that E is positively homogeneous, i.e. E(aqp) (E) any μ E P a 2 0. Then for all p E [1, oo), E is not bi-Lipschitz with respect to Wp-
llw, defined in (39) for distributions in P&N (IRd). Note that this pseudonorm is positively homogeneous, namely, for all ja PSNN (RR.). 1) a 2 0 and
Proof. In our proof we use the pseudonorm ||
p E [1,
|aullyy w, = allulll,,
Let ) be sequence of real numbers such that
< 0c+1 S 1 S
VI 2 1.
(53)
By assumption, the set Ω contains a ball (0) of radius r > 0 centered at zero. Choose x + 0 in that ball. For 0 [0, 1], define the distribution
(()) (1 0)6o + 08
where dg is Dirac's delta function at E Rd. That is, n() assigns weights of 1 0 and # the points 0 and æ in Id respectively. It is straightforward to verify that for 1 S p < cc
Wp(n(O)), 50) = [0.1| (0.1|=1|0)\//p =V V0
This holds for p = oo too, with the convention Ve 1 for x > Therefore, for all e N, E(u(0.)) E(6o)
E(H(O.)) - E(do) [
E(p(O:))
(54)
(p(p())),
d0,
| 70,
where the last equality follows from the homogeneity of E, which implies that E(5v) 0.
We can assume that E in upper-Lipschitz, since otherwise there is nothing to prove. Under this assumption, the norm of the expression in (54) is uniformly upper-bounded for all I. E N, which implies that there exists a subsequence of 0, for which this expression converges. Replacing 0, with this subsequence, it can be verified that the new subsequence still satisfies (53), and that for an appropriate vector L E R
lim E(u(O.) L. t000
d0
Now, consider the sequence distributions
0:
N4 ==
V0
0
u(0,-1), 1 122.
29

---

Published as a conference paper at ICLR 2025
Since
I 81-1 V <1
-
and x E (0) C Ω the distribution Hu indeed belongs to PsN (22). We wish to lower-bound the p- Wasserstein distance from pu(Ot) to ju for 2. Note that both distributions divide their mass between zero and an additional vector: the distribution jr assigns a mass of 0z-1 to the nonzero point V " whereas the other distribution μ(0,)) assigns smaller mass of 0, to the nonzero point æ, with the complementary masses assigned to zero. Therefore, transporting ee to ((0e)
requires transporting at least mass of 01-1 02 from l 0, that for all S p oo
(T(((()), (i) (0,-1 01)1
0.
0
61-
0,(1
||
0,-1
1 0el
1
0e||
We obtained that for p oc,
(o))) 2 //
(55)
and the same argument can be used to verify that this is the case for = oo as well. We conclude that
E(u(O.)
V V B(u(0r11))
B(n(0))- - E(POll
Th
(,(p(0)). (.)) D. B(u(O.)) 8. E(n(Or-1))
2/1/2|1011
+ 0
1//2||=
where (a) is by (55) and the homogeniety of F, and the convergence to zero is because both ex- pressions in the numerator converge to the same limit L. This shows that E is not lower-Lipschitz, which concludes the proof of Lemma 10
The following lemma shows that the homogeneity assumption on E can be released.
Lemma B.11. Let E PsN (2) - Rm, with Ω C I&d being an open ball centered at zero, N and > 1. Then for all oo). E is not bi-Lichchitz with respect Wp.
Proof. Let E [1, oo] and suppose by contradiction that E is bi-Lipschitz with constants 0 <<c C. < xi,
c. Wp(H < (x) - E(Ail < C (p A))
E
(56)
We can assume without loss of generality that E(6o) = 0, since otherwise, let
i(w) = E(u) F(5o).
then E satisfies (56) if and only if f satisfies (56).
We first prove an auxiliary claim.
Claim. For any μ. js E PSN (92) with ||ll],, = and 0 < llAll,, 1
E
E(a)
SC. -C-(1-1121))) 1-I1ūllw, SC CW,(u,
(57)
77 1ww,
30

---

Published as a conference paper at ICLR 2025
Proof. By (56).
- E(H)
Tiullw )
1
E
3 SC. w'
lilllwo
We shall now show that
p
Wp
10
S 1 11-llälll
Let = N w be parametrization of j.. Consider the transport plan (TT)),,EE()) from
a given by
TAll,
W/ i= j 10
Tij
i i j.
By definition, W
B TAlli
is smaller or equal to the cost of transporting j to TTll,
according
to æ. Thus, for p < 0c.
N S
N
| TTalio -1 $t. P
)
W
TAll
Wi
jillwp
W/
1=1
1=1
1
Tallop -1) P ll:l
Tpllww
Σ
w/
|2,11"
1=1
1
llūll
and thus
llallw.) 1 S (1-17m
W
lill,,
Both sides of the above inequality are continuous in p. including at the limit p - co. Thus, the above inequality also holds for p co. Now, to show that
l1,ll,, S Wp(u,p)
note that
[,l,,, IIul|w, lllll, Wo(u.0) W)(2,0) S (,,,,,),
where the last inequality is the reverse triangle inequality, since " ) is metric. Thus, (57) holds.
Now we define the homogenized function E (S.) - 1m++ by
B(u) = [Iellww, 11alb,, E ((AltT)) -
+ do
(58)
/ d0-
It is straightforward to verify that É is well-defined and positively homogencous. By Lemma B.10, È it is not bi- bi-Lipschitz with respect to Wp. and thus there exist two sequences of distributions Me E (P) 1, such that
B(m)) - E(tt))
L,
(59)
Wy(ms, (i))
t-000
with L = or L = co. Since É is positively homogeneous, we can assume without loss of generality that
1 ||lellww,
I1Aellw, for all t.
31

---

Published as a conference paper at ICLR 2025
This can be seen by dividing each He and Re by max { for all for which I|elll,, < <IlAe llw,
[Ileellw, |1 Selll,, } and swapping μr and p
jr = 60 for an infinite subset of indices 6, then redefine pe and ir to be the corresponding subsequences with those indices, and now (59) implies that
[E(w) E(do)l (a) B(wx) E((b)
L,
Wip 60)
(1)
1-00
with (a) holding since lIAellw; = 1. This contradicts the bi-Lipschitzness of E. Thus, we can assume that jāe- = do at most at finite subset of indices t. By skipping those indices in μr and jús we can assume without loss of generality that
= |eellw,
|| w, > 0 for all t 2 1.
(60)
Let us first handle the case co. The first coordinate of B(He) E(pe) is bounded by
|Ilellw; - IlAellw, = - l1Alll,, S Wp(ue, pe)
according to (57). Thus, the violation of bi-I ipschiichssss must come from the last mn coordinates of É(w)) - (((a). That is, (59) with L = oo, combined with the fact that [| > 0 Vt, implies that
IIleellww, F
IImellww )
Ms
I1aellww, E
Dr |Iiallww
O0.
(61)
Wp
On the other hand,
||pellwe ) 11 E
|1mellow E
( E(He) l1@elllw E
Pre
|I
|1 ll llwi
(b) |E(()) E(Pill +
jut (Tuellw,
Tiee
E((e)) - E3
+
E3
E.
Li
| [[iellw,
(62)
where (a) holds since [|æellww, and (b) is by the triangle inequality. We shall now upper-bound the three above terms.
First, by (56).
[ - B(Te)l < C- Re).
(63)
Second,
fuu 1
F(Oe) F. F
< SC. (CWppuu..))
(64)
by (57). Lastly,
Tee E
)
fit l
)
Fi
-1ieellw E
E
0
lipellws,)
lliellls
(1 (1Aeelbos)
1
E(oo)
I1ūellw,
(1 (1- l1
C. Wr
The
80
(65)
1Tieellwsp
1 [1 c.
I WW
C. (1 [1mlll
(6) C. C-Wpppe pc).
5
32

---

Published as a conference paper at ICLR 2025
where (a) is by (56), (b) is by the definition of || Illy, and (c) is by (57). Incorporating (63)-(65)
into (62) yields
Iivelll,, E
pe [1reellwwe/
) - 11 11ielllw, E
joe |/
)
S3C 3C (p():
which contradicts (61).
Let us now handle the case L O. For two sequences of numbers ar R, t 2 1, we say that
0r = o(b))
at lim 1-000 by
0.
Denote
d = Wp(pe,me)
According (59) with L 0, the first component of E(ua) B(ūz), which equals || Il Iillll,, satisfies
Iluellw, - 1lūtellw O,
Wplmes Dea)
10000
and thus
11 |1lellw,, -1ieellw, = o(d)).
(66)
By the triangle inequality,
E(()) E(Po)ll 5
The
The
E((e) - 11i E'
lTelll
+
llw,
E'
ITtellws
i ww
(67)
+
E
E(its))
|
Ilw,
We shall show that each of the three above terms is o(d4).
First, since [|l 1,
| llw, )
The
E(Hu) - liellww, E
(68)
|llaelbw, E
ITwellww,
- l1äellw, E
|1
)
(8)
<
D(pe) - E(a+))
(59) o(de),
33

---

Published as a conference paper at ICLR 2025
with inequality (a) holding since the argument of the norm the left consists the last coordinates of E((4) B(pe). Thus, the first term in (67) is o(de). For the second term,
ju
j w; E
E
l|ãelll,
|pellw,
l|ātelww))
E
l liolllwp
The
E.
(d(o))
limell
(69)
C. w
ihe Ilãeellw
1 (
50
l
c.
|1 l1ie,ll,,
(b)
o(de).
where (a) by (56) and (b) is by (66). Finally, by (57).
E
T
)
E(ūe)
SC. (1 [[ūellw)) o((d))
(70)
[|@ellw,,
Therefore, by (68)-(70) and (67). we have that
[[E(me) E(pe)ll o(d)),
and thus E is not lower-Lipschitz. This concludes the proof of Lemma B.11.
To finish the proof of Theorem 4.4, suppose that C IRd an arbitrary set with a nonempty interior.
Let 2o Ω be open ball contained in S2, and let 2% be the center of Ω4g Then 2g 2g open ball centered at zero.
Given E: PgN (i) I define : ::PN(Sa 2p) IR by
E(u) = E(u + 2o).
where μ + æo the distribution μ shifted by æb:
N
1 + Zo = Σ
Then E satisfies the assumptions of Lemma B. and thus there exist two sequences of distributions Hs Be E (fo - z0) such that
E(p,) - E(A+)
L,
with L = 0 or co. Note that the sequences of shifted distributions {ur + xu and (it + 0ohe21 in P (P,) and thus in (S)) Since
Wplue + Zo. Tee x0) Wplue,
have that
E(u + xo) E(ūa æo)1
||E(e + xo) E(ūt + xo) Wp Fi))
Wp(m4 + Lu. fi + xo)
[((He) E(He)
1+000 T.
Wp Pt)
which implies that E is not bi-Lipschitz PSN (Ro), and thus not Ps)
34

---

Published as a conference paper at ICLR 2025
We now prove Theorem 4.2. Before proceeding with the proof, we first recall the precise definition of piecewise linearity.
Definition B.12. A function f IRm - R" is piecewise linear if there exists a finite collection of polyhedra 8 {P.) that partition Rd (i.e., their union covers Ih and their interiors are disjoint)
such that f is linear on each P.. That is, for each i, there exist a matrix A, E I and a vector b, E R= such that
f(z) = A, + bi, for all x E P..
Theorem 4.2. (Statement in Section 4) Let E' : PgN + -h be injective and positively homogeneous. Let w. in E A be fixed, and suppose that the functions E(X, w) and E(X, .) are piecewise linear in X. Then there exist C, 0 Such that for all x X E Edd and [o).
c-Wo (XX, w) (x, x)) < |((X w) E X 2 || S.C (Xx w)) [x, 20) ) (12)
E(X,
(x
Proof. The proof is outlined as follows: First we show that there exist constants č, C > 0 for which (12) holds in the special case p Then we show that for any fixed w, t E A N there exists constant 3 > 0 such that for all X, X RdNN
Wi (Xx w) X 2n 2 B- B- (ox((X, w)), (x, (x,t))) 10)
(71)
This will imply that for the given pair w, w. (12) holds for all p E [1, oo] with the constants c= c= cB and C = C. To see this, note that by the monotonicity of W; with respect to p (Eq. (52)),
Wi (1((,,u) w) :)) - Wp (X, (x) (x, S Woc (x, w), (x, i
(72)
Therefore, Eq. (12) with p and constants Č, and Equations (71) and (72) imply that [[(X w) E X )
ス ē. Wi (X. w.. X の
2 28 Woo (x) w). (x, w))
(73)
2 ē. Wp (X,w), (X,ú) (
and ((X, w) E X wb) S wi (X, w)), (x, (xx,0)) SBB W/ ((x, w) (x, (74)
Let us begin by proving that (12) holds for p = The I-Wasserstein distance between two distri- butions parametrized by (X, w) and (X. w can be expressed by
Wi (XX w) (x, )0))
min w
Σ
Tig ( 5601
(75)
[JE[NN
where the set 11(,, w) of transport plans from (X, w) to (x, w is given by
N Σ nq mmz= wi AVjEIN [N Σ Tyj
N
11(w, w) [0, NxN Vi [N]
ij
In particular, Ilw w depends only on w and v and not on the points X, X.
Let W1 be a modified I-Wasserstein distance that uses the fn-nrrm rather than Eg as its basic cost function:
Mi ((X, (X, (x,i)) = min =
Σ Tig ( a00
(76)
[JE[N]
Note that since
| < ||:|||| S Vdll0112
Va E Rd
(77)
A polyhedron is finite intersection of closed half-spaces.
35

---

Published as a conference paper at ICLR 2025
have Wi (X), (x,z)) 2)) - Wi (X, w). (x, () s vd. wi 22((2 (X, w). (x,²)):
(78)
Thus, it suffices to prove an analogue of (12) in which W1 is replaced by Wi and || by 11:11 namely, that for fixed w, w. there exist constants C, such that
Wi (X, (X,,), w) [x, ū - ww) E X i |l, C. Wi (x) (x, (79)
To this end, we define the function f RaxN BaxN - IR.2
f x = || w) w) - EE ( i
Wi (i((X, w)) X |b
Since both E and Wi are positively homogeneous with respect to
[ x)) (xx))
then so is f. Namely,
(ax,ax) = xx,,
Va 2 0.
(80)
In addition to being positively homogencous, is also piecewise linear, as we will show below. We then prove that these two properties together imply that E is bi-Lipschitz with respect to W1.
Let us now show that f is piecewise linear in x
The first component of f.
() - EXX, i || is clearly piecewise linear, as it is the composition of the -n-mrmm with piccewise-lincee function. We shall now show that the second component Wi ((X, w) X w
is also piecewise linear. For any fixed and X, the optimization problem in (76) is a linear pro- gram in r,, with the set of feasible solutions being the compact polyhedron 11(w, w... Thus. the optimal solution must be attained at one of the vertices of II w, w). As any polyhedron has a finite number vertices", let x(1), ('))..., ((K) be the vertices of II(w (w, w) and recall that these vertices do
depend on [X x). Therefore, (76) can be reformulated
Wi ((X, w). (x, 1 )) = min
Σ
(8) x())
x
(81)
Ti
KE(K
[E([N]
From (81) can be seen that Wi w)), X t is piecewise linear (x, x). as the minimum of finite number of piecewise-linear functions. Since the concatenation of piecewise- linear functions is also piecewise linear, we have that f fx, is piecewise linear.
Now, let A C IR2 be the image of f: A == (1(,xx
Since f is piecewise linear, it maps the space Rdx N x Rd to finite union of closed polyhedra (some of which may be unbounded). Hence, A is finite union of closed sets, and thus is closed.
Denote the two coordinates of f by
fi X, x = |E(X w) EE (EXXX,,) l1:
fz x, x) = Wi ((X, (X, 2
Since E is injective by assumption, if fi ( = then fo [X X = 0. Thus, A does not contain the point (0, 1). Conversely, if fz X, X) same distribution, that is
= O, then (X,w) and (x i represent the
w (4)
Σ
See (Grünhaum, 2003), Theorem 3, page 32, and the definition of polyhedral on page 26 therein.
36

---

Published as a conference paper at ICLR 2025
and since E(X, w) PsN (R*) + R depends only the input distribution and not its particular representation, have that fi [x X
O. Hence, A does not contain the point (1, Now, suppose by contradiction that there do not exist constants C, > 0 for which (79) holds. Then, there exist two sequences X,, X, E RdxN such that both coordinates / X X. are nonzero for all E N, and their ratio satisfies
fi (x, - I,
(82)
f2 X1, x1: X:)
where the limit L is either 30.
Suppose that L = cc. By the positive homogeneity of E and W1, we can divide both X2 and X, by Si Xt, X.) while still satisfying (82) with L = xc. With these new sequences, we obtain
fi X1, X,
VEN
fa X X,
0.
It follows that there exists sequence in A that converges to the point (1,0). Since A is closed, conclude that (1 E A, contradicting our assumption.
By similar argument, if L O, then (0, 1) A, leading to the same contradiction.
To finish the proof, it left to show that (71) holds with some constant B >> 0 assuming that w and i are fixed. To this end, define the sets 1& C [N]² for A: E [K], with K as in (71),
I& (1.5)) [N]
(1)
m
By the definition of I[, w) for all l E [K], Is is nonempty. Let
= min (K5) Ti
(K)
h [K]
(83)
That is, cg the smallest positive entry of (k)
and by the above, indeed c Let
Gmin
min Cx > 0.
(84)
K
Thus, for all X, X E IRdx
[ (, )) (78) W (X, w)
Vd. Wi (X,
(X, 20)
min
(K)
(1) x(5)
min
Tie (x)))
(k)
kE[K]
[JENN
[JEN[N]
(k)
x001 50
min Σ (1.5)
(0 50
min
Tij
kE
kE(K)
(1))e4
(84)
x0))
x00)
2010 x0,1
min
Emin
> min
max (T.
Emin
kE K]
K
(1.))1
20(1)) x0
ij E [N], T > 0
(R)
Emin
min max kE[K]
max { | -x(i)
ij E [N], Tij > 0}
Emin
min xE(T Jm
Emin
min =EIT(a0,00)
max { x | 2000111
-
[N], mij >0}
Wac (X, w), (x,
Emin
37

---

Published as a conference paper at ICLR 2025
where (a) is since
(k) 0 whenever (i,5) Ixi (b) is by the definition of Ixi (c) is simple iij
reformulation; (d) is since the minimum is taken over larger set I[( w) 2 (x) [==1'
[K]
and (c) is
by the definition of Woo Hence, (71) holds with 3 fmin and the theorem is proven.
vd
38

---

Published as a conference paper at ICLR 2025
C
EXPERIMENT DETAILS
Hardware
All experiments were conducted on a single Nvidia A40 GPU.
C.1 COMPARISON WITH PSWE
In our comparison with PSWE, shown in Table 1, we used pairs of multisets of the form X
for two different values of n, denoted nu n2, where 1 E R is the vector whose entries
all equal For example, one could choose
X1 ()
-
Despite the fact that these multisets distinct, the interpolation used by PSWE (sce c.g., Naderial- izadeh et al. (2021)) identifies these two multisets with the same continuous distribution, and hence assigns them the same embedding. In Table we used d = 3 and n1 5 and 72 = 200, and found, as expected, that PSWE assigns the same value to both multisets, up to numerical error, regardless method successfully separates X1 and X₂
of the embedding dimension m. In contrast, with m 1. and yields a good approximation of the sliced Wasserstein distance as m increases. Similar results to those shown in the table were obtained for other choices of d, 711, n12-
Note that computing the sliced Wasserstein distance between such pairs is straightforward, since, for any two multisets 11, 2 in R2 with all points belonging to the same line, can be shown that
SW(X,, X2)
W ,(X,, X2),
Vd
and W(X), X2) can be computed using a linear programming solver.
The PSWE method takes two input parameters- the number of slices L and embedding dimension N, which, unlike our method, need not be identical. We set both parameters to the embedding dimension
C.2 LEARNING TO APPROXIMATE THE WASSERSTEIN DISTANCE
In this experiment we used embedding dimensions m. = me 1000. The MLP consisted of three layers with hidden dimension of 1000. With this choice of hyperparameters, our model has roughly million learnable parameters and 5 million parameters in total. These hyperparameters were picked manually. The performance of architecture did not exhibit high sensitivity to the choice of hyperparameters: on most datasets, similar results were obtained with MLPs consisting of 2 to 8 layers, and with hidden dimensions of 500, 1000, 2000 and 4000.
We used fixed parameters for the first embedding Ej and learnable parameters for the second em- bedding Ez. This choice was made since E₁ is, most cases, supposed to handle arbitrary input point clouds, whereas the input to Eg is more specific, in that is always a set of two vectors that are outputs of Ej. Thus, in principle the architecture may benefit from tuning Eg its particular input structure. In practice, using fixed parameters in both embeddings did not significantly impair performance.
Remarkably, applying an MLP to the input points prior to embedding them via Fg (i.e. adding feature transform), as well as applying an MLP to the two outputs of E₁ prior embedding them via Es, impaired rather than improved the performance. This indicates that embedding is expressive enough to encode all the required information from the input multisets in a way that facilitates processing by the MLP क. thus making additional processing at intermediate steps unnecessary.
Inference times for one pair of multisets were less than half a second for the Mode Noee -are-e dataset, and less than 0.2 seconds for the rest of the datasets. The training times of the competing models appear in Table 3.
Training performed on Nvidia A40 GPU, whereas the rest of the methods were trained over an Nvidia RTX A6000 GPU, both which have similar performance on 32-bit floating point (37.4 and 38.7 TFLOPS).
Exact computation of the I-Wasserstein distance using the ot. emd2 () function of the Python Optimal Transport package (Flamary et al., 2021) was up 2.5 times slower than our method (2 to
39

---

Published as a conference paper at ICLR 2025
5 ms 1.9 ms) small multisets (less than 300 elements) and 150 times slower (640 ms vs 4.2 ms) on large multisets [ModelNets -large))))))...
C.3 ROBUSTNESS TO PARAMETER REDUCTION
In both architectures, the input and output dimension of each MLP, well as that of the FSW embedding, were multiplied by scaling factor to obtain the desired number of parameters. Training our model all problem instances took between 60 to 65 minutes. Training PointNet took between 4:43 hours to 5 hours, and was done using the original code of (Xia, 2019).
40


Reviews:
[{"review_id": "GacppfgBCr", "reviewer": "Reviewer_Ew1o", "summary": "This paper considers Fourier slicing embedding both for a collection of probability distributions and multisets over $\\mathbb{R}^d$ and supported at $n$ points. The embedding consists of a projection sample on a 1-dimensional vector on the sphere then calculates a cosine transform of the projected quantile function. Under a specific probability distribution of the frequency, the authors prove that the expectation of the estimation error between the embedded measures is exactly the sliced Wasserstein distance. A second part of the theoretical results consists of proving the injectivity of the embedding under the assumption that the dimension embedding $m \\geq 2n(d+1) +1$. Numerical experiments are conducted on point cloud classification.", "questions": "See Weaknes section.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- The paper is well-written and easy to follow. Proofs are rigorous. \n- Proposing the sliced embedding Wasserstein (SEW) through a cosine transform of the projected quantile function. Sampling the quantile function via cosine transform is novel.\n- Injectivity and bi-Lipschitz properties of FSEW on the collection of multisets.\n- Numerical experiments showcase better Wasserstein approximation on simulated datasets and three real datasets than NProductNet, WPCE, NSDeepSets, and Sinkhorn.", "weaknesses": "- Several approaches for the derivative of sliced Wasserstein distance like, distributional sliced Wasserstein (Nguen et al, ICLR'21), max-sliced Wasserstein, etc ... Could you highlight the difference between FSW and the SOTA derivative of sliced Wasserstein?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "6Y7MWhkAt4", "reviewer": "Reviewer_kJBi", "summary": "This paper presents a novel approach to high-dimensional dataset embedding. The authors provided theoretical performance guarantees and numerical study to show the superior performance of the framework.", "questions": "I am new to this field. Could the authors elaborate more on the practical motivation and applications of this approach?", "limitations": "", "rating": 5, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The theoretical contribution seems to be sound, with explicitly stated technical assumptions and results. Numerical study is solid.", "weaknesses": "1. The authors denovted much space to describe the p-wasserstein and infinity-type Wasserstein distance. Why it is necessary to introduce infinity-type Wasserstein distance?\n2. In line 222, the authors mentioned that in the special case of d=1, Wasserstien can be computed significantly fast. So what is the complexity rate?\n3. In line 344, what is the definition of STD???\n4. The authors should provide proof ideas for the main technical results in the main content.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "PVCP4i5tks", "reviewer": "Reviewer_gmod", "summary": "The paper seeks to establish a mapping from multisets and measures over $ \\mathbb{R}^d $ into Euclidean space, ensuring that the sliced Wasserstein distance corresponds to the distance between their mappings in the target space. The authors propose a mapping that is bi-Lipschitz for multisets and injective for measures. Additionally, they demonstrate that a bi-Lipschitz map for measures does not exist.", "questions": "I haven't any question.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "The paper is well-structured, and its message is clear. The proofs provided are rigorous and exceptionally clear. This particular problem is quite interesting. I really enjoyed reading the paper.", "weaknesses": "I don't see any weaknesses. Therefore, I recommend accepting it.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "nkyzy6KHJZ", "reviewer": "Reviewer_styv", "summary": "**Summary:**  \nThe paper introduces the \"Fourier Sliced Wasserstein (FSW) embedding\" for data in \\(\\mathbb{R}^d\\).\n\n**Theoretical Contributions:**  \n1. The authors prove that the embedding preserves or approximates the sliced Wasserstein distance.  \n2. They also demonstrate that the embedding technique is injective and bi-Lipschitz.\n\n**Numerical Experiments:**  \n1. The authors evaluate the approximation error of the proposed Fourier Sliced Wasserstein embedding.  \n2. They showcase an application of FSW for approximating the Wasserstein distance using a Multi-Layer Perceptron (MLP).", "questions": "1. Regarding point (2), is \"Multisets\" simply another term for \"discrete distributions\"?  \n2. Could you clarify what \"E2\" refers to in lines 473-474?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. The combination of the Fourier/cosine transform and the sliced Wasserstein distance (see Eq. (6)) is a novel approach.\n2. Theoretical properties for this new technique with respect to the uniform distribution, along with its empirical approximation, are proposed (see Theorem 3.2, Corollary 3.3).\n3. Injectivity and bi-Lipschitz properties of the embedding have been investigated.", "weaknesses": "1. I recommend adding a section to introduce baseline methods. For example, explaining how Sinkhorn [Cuturi, 2013] can be used to train a neural network as a Wasserstein distance estimator. Currently, the experimental setup (E1, E2, Phi, Leaky-ReLU) appears tailored only to the proposed method in this paper.\n2. It would be beneficial to introduce a real-data application of the proposed Sliced Wasserstein distance embedding technique to illustrate its practical utility.\n3. I’m unclear on why 'bi-Lipschitz' is considered a crucial property. Could you provide an example to clarify? For instance, in which applications would the lack of a bi-Lipschitz property cause issues, and where having this property could offer distinct advantages?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)