Neural approximation of Wasserstein distance via a universal architecture for symmetric and factorwise group invariant functions
Authors: Samantha Chen
Yusu Wang
Venue: NeurIPS 2023 poster (2023)
Keywords: neural networks
Wasserstein distance
universal approximation
optimal transport

Abstract: Learning distance functions between complex objects, such as the Wasserstein distance to compare point sets, is a common goal in machine learning applications. However, functions on such complex objects (e.g., point sets and graphs) are often required to be invariant to a wide variety of group actions e.g. permutation or rigid transformation. 
Therefore, continuous and symmetric *product* functions (such as distance functions) on such complex objects must also be invariant to the *product* of such group actions. 
We call these functions symmetric and factor-wise group invariant functions (or SGFI functions} in short).
In this paper, we first present a general neural network architecture for approximating SFGI functions. The main contribution of this paper combines this general NN with a sketching idea in order to develop a specific and efficient neural network which can approximate the $p$-th Wasserstein distance between point sets.
Very importantly, the required model complexity is *independent* of the sizes of input point sets. 
On the theoretical front, to the best of our knowledge, this is the first result showing that there exists a neural network with the capacity to approximate Wasserstein distance with bounded model complexity. Our work provides an interesting integration of sketching ideas for geometric problems with universal approximation of symmetric functions. 
On the empirical front, we present a range of results showing that our newly proposed neural network architecture performs comparatively or better than other models (including a SOTA Siamese Autoencoder based approach). In particular, our NN generalizes significantly better and trains much faster than the SOTA Siamese AE.
Finally, this line of investigation could be useful in exploring effective neural network design for solving a broad range of geometric optimization problems (e.g., $k$-means in a metric space).

Neural approximation of Wasserstein distance via a universal architecture for symmetric and factorwise group invariant functions
Samantha Chen
Yusu Wang
Department of Computer Science Engineering University of California San Diego sac003eucsd.edt
Halıcroğuu Data Science Institute University of Californi San Diego
yusuwangeucsd.c edu
Abstract
Learning distance functions between complex objects. such as the Wasserstein distance to compare point sets, is common goal in machine learning applications. However, functions on such complex objects (e.g., point sets and graphs) are often required to be invariant to wide variety of group actions e.g. permutation or rigid transformation. Therefore, continuous and symmetric product functions (such as distance functions) on such complex objects must also be invariant to the product such group actions. We call these functions symmetric and factor-wise group invariant functions (or SFGI functions in short). In this paper, we first present a general neural network architecture for approximating SFGI functions. The main contribution of this paper combines this general neural network with sketching idea to develop a specific and efficient neural network which can approximate the p-th Wasserstein distance between point sets. Very importantly, the required model complexity is independent of the sizes of input point sets. On the theoretical front, the best of our knowledge, this is the first result showing that there exists neural network with the capacity to approximate Wasserstein distance with bounded model complexity. Our work provides an interesting integration of sketching ideas for geometric problems with universal approximation of symmetric functions. On the empirical front, present range of results showing that newly proposed neural network architecture performs comparatively or better than other models (including SOTA Siamese Autoencoder based approach). In particular, our neural network generalizes significantly better and trains much faster than the SOTA Siamese AE Finally, this line of investigation could be useful in exploring effective neural network design for solving broad range of geometric optimization problems (e.g., k-means a metric space).
1 Introduction
Recently, significant interest in geometric deep learning has led to a focus on neural network
architectures which learn functions complex objects such point clouds [4 2 and graphs [23)
Advancements in the development of neural networks for complex objects has led to progress in
variety of applications from 3D image segmentation (28) to drug discovery 3 [3] One challenge
in learning functions over such complex objects is that the desired functions are often required to
be invariant to certain group actions. For instance, functions on point clouds often permutation invariant with respect to the ordering of individual points. Indeed, developing permutation invariant or equivariant neural network architectures, as well as understanding their universal approximation properties, has attracted significant attention in the past few years; see e.g., 22
T9
20
25
16
33 29 30
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

However, in many geometric or graph optimization problems, our input goes beyond a single complex object, but multiple complex objects. For example, the p- p-Wasserstein distance W, (x Y) between two point sets X and Y sampled from some metric space (e.g.. IRd) is a function over pairs of point sets. To give another example, the 1-median of the collection of ls point sets Pi...., PB in R can be viewed as a function over k point sets.
A natural way to model such functions is to use product space. In particular, let X' denote the space of finite point sets from a bounded region in Rd. Then the p- p-Wasserstein distance can be viewed as function Wp p X' x X' - IR. Similarly, 1-median for k point sets can be modeled by a function from the product of & copies of X to R. Such functions are not only invariant to permutations of the factors of the product space (i.e. Wp (X, Y) = Wp Wp(Y, X)) but are also invariant or equivariant with respect to certain group actions for each factor. For p- W--ssssssttinn distance, W, is invariant permutations of the ordering of points within both X and Y. This motivates us to extend the setting of learning continuous group invariant functions to learning continuous functions over product spaces which are both invariant to some product of group actions and symmetric. More precisely, we consider type of function which we denote as an SFGI product functions.
Definition 1.1 (SFGI product function) Given compact metric spaces (K;, dx.) where i E [k], we and
define symmetric and factor-wise group invariant (SFGI) product functions as f X2 x X2 x: Xx - R where I is symmetric over the le factors, and (2) invariant to the group action G1 x G2 .......... Gk for some group G3 acting on Xy for each i [1, k..
Contributions. SFGI product functions can represent a wide array of geometric matching problems including computing the Wasserstein or Hausdorff distance between point sets. In this paper, we first provide general framework for approximating SFGI product functions in Section 3.1 Our primary contribution, described in Section 3.2 is the integration of this general framework with a sketching idea in order to develop an efficient and specific SFGI neural network which can approximate the p- Wasserstein distance between point sets (sampled from a compact set in a nice metric space, such as the fixed-diment died-diienioonal Euclidean space). Most importantly, the complexity of our neural network (i.e., number of parameters needed) is independent of the maximum size of the input point sets, and depends on only the additive approximation error. To the best of our knowledge, this is the first architecture which provably achieves this property. This is in contrast to many existing universal approximation results on graphs or sets (e.g., for DeepSet) where the network sizes depend on the size of each graph or point set in order to achieve universality (20 30 4]. We also provide a range of experimental results in Section 4 showing the utility of our neural network architecture for approximating Wasserstein distances. We compare our network with both a SOTA Siamese autoencoder [15], a natural Siamese DeepSets network, and the standard Sinkhorn approximation of Wasserstein distance. Our results show that our universal neural network architecture produces Wasserstein approximations which are better than the Siamese DeepSets network, comparable the SOTA Siamese autoencoder and generalize much better than both to input point sets of sizes which are unseen at training time. Furthermore, we show that our approximation (at inference time) is much faster than the standard Sinkhorn approximation of the p- Wassersstnii distance at similar error threshholds. Moreover, our neural network trains much faster than the SOTA Siamese autoencoder. Overall, network is able to achieve equally accurate or better Wasserstein approximations which generalize better to point sets of unseen size as compared to SOTA while significantly reducing training time. In Appendix B we provide discussion of issues with other natural choices of neural network architectures one might use for estimating Wasserstein distances, including one directly based on Siamese networks, which are often used for metric learning.
All missing proofs/ details can be found in the Supplementary materials.
Related work. Efficient approximations of Wasserstein distance via neural networks are an active area of research. One approach is to use input convex neural networks to approximate the 2- Wasserstein distance 18 22]. However, for this approach, training is done per pair of inputs and restricted to the 2- Wasserstein distance which makes it unsuitable for general neural network approximation of Wasserstein distances between discrete distributions. This neural approximation method contrasts with our goal: a general neural network that can approximate the p- Wasserstein distance between any two point sets in compact metric space to within e-accuracy. Siamese networks are another approach for popular approach for learning Wasserstein distances. Typically, a Siamese network is composed of single neural network which maps two input instances to Euclidean space.

---

The output of the network is represented then by f--norm between the output embeddings. In [S]. the authors utilize Siamese autoencoder which takes two histograms (images) as input. For their architecture, single encoder network is utilized to map each histogram to an embedding in Euclidean space, while a decoder network maps each Euclidean embedding back to an output histogram. The Kullback-Liebler (KI.) divergence between original histogram and the output histogram (i.e. reconstruction loss) is used during training to regularize the embeddings and the final estimate of 2- 2-Wasserstein distance is the lg norm between the embeddings. The idea of learning Wasserstein distances via Siamese autoencoders was extended in [5] to point cloud data with the Wasserstein point cloud embedding network (WPCE) where the original KL reconstruction loss was replaced with a differentiable Wasserstein approximation between the original point set and a fixed-size output point set from the decoder network. In our subsequent experiments, we show that our neural network trains much more efficiently and generalizes much better than WPCE to point sets of unseen size.
Moreover, the concept of group invariant networks was previously investigated in several works, For instance, DeepSets [4 and PointNet [[9] are two popular permutation invariant neural network which were shown to be universal with respect to set functions. In addition to group invariance, there have also been efforts to explore the notion of invariance with respect to combinations of groups, such as invariance to both SE(3) and permutation group 21 or combining basis invariance with permutation invariance [17. Our work differs from previous work in that we address universal neural network which is invariant with respect to a specific combination of permutation groups that corresponds to an SFGI function on point sets. In general. we can view this subgroup of the permutation group encoding the invariance of each individual point set with the symmetric requirement of the product function corresponds to a specific subgroup of the permutation group. Thus, previous results regarding permutation invariant architectures such as DeepSets, PointNet or combinations of group actions (such as [11] do not address our setting of SFGI functions or p- --Waseerttein distances.
including 44 22
2 Preliminaries
We will begin with basic background on groups, universal approximation, and Wasserstein distances.
Groups and group actions A group G is an algebraic structure that consists of set of elements and a binary operation that satisfies a specific set of axioms: (1) the associative property, (2) the existence of identity element, and (3) existence of inverses for each element in the set. Given metric space d,), the action of the group on X is a function a : G x X' + X that transforms the elements of X for each element E G. For each element π E G. we will write
to denote the action of a group element π on x instead of (1 x). For example, if G is the permutation group N). and X R N, then for any E G, π represents the permutation of elements in œ E R N via IT i.c given x = (I1, T2,..., π x = (I+(1); (2()) =z(N)) A function f X² - y is G-invariant if for any X E X' and any x E G, we have that fXX = f(x -X).
over [N] {1, 2,
Universal Approximation. Let C(X', R) denote the set of continuous functions from metric space (X, dxx to R. Given two families of functions Fi and F₂ where F, S F2 and F1, F2 S C(X, R). say that Fi universally approximates F2 if for any > and any E F2, there is a g E Fi such that ||g - flloe < c. Different norms on the space of functions can be used, but we will use Lo norm in this paper, which intuitively leads to additive pointwise-error functions. The classic universal approximation theorem for multilayer perceptrons (MLPs) 7 states that feedforward neural network with a single hidden layer, using certain activation functions, can approximate any continuous function to within an arbitrary additive e-error.
the domain of these
Permutation invariant neural networks for point cloud data One of the most popular permuta- tion invariant neural network models is the DeepSets model defined in (34. DeepSets is designed to handle unordered input point sets by first applying a neural network to each individual element, then using sum-pooling to generate an embedding for the input data set, and finally, applying a final neural network architecture to the ouput embedding. Formally, suppose we are given finite multiset S {x x E HR (meaning that an element can appear multiple times, and the number of times an
3

---

element occurs in S is called its multiplicity). The DeepSets model is defined as
NirepSec (S) 90a
Σ DES
ho, (1:)
where he, and g0, are neural network architectures. DeepSets can handle input point sets of variable sizes. It was also shown to be universal with respect to continuous multiset functions.
Theorem 2.1 ((44) 12) Assume the elements are from a compact set in IR*, and the input multiset size is fixed as N. Let = 2kN + 1. Then any continuous multiset function, represented as
f: R + which is invariant with respect to permutations of the columns, can be approximated
arbitrarily close in the form of p Emexxxi b(x)) for continuous transformations & : IR + IR and
p: IR. R
While universality for the case when &o 1 was shown using symmetric polynomials, the case for k > 1 in fact is quite subtle and the proof in (34) misses key details. For completeness, we provide a full proof in Appendix E.1 for when the output dimension of á is 1 = ( (*++)), It was recently shown in [22 | that the output dimension & can be reduced 2kN + 1, which is the dimension of which we use in Theorem 2.1 and subsequent theorems. In both the cases where the output dimension of á is = (*+-)) or 2kN + 1, Theorem 2.1 implies that if we want to achieve universality, the required network size depends on input point cloud size.
Wasserstein distances and approximations. Here we will introduce Wasserstein distance for discrete measures. Let (X, dx) ) be a metric space. For two weighted point sets P {(z,, w.) az e x, E. i, [n]] and distance between P and Q as
[(1%, w!) 2 X. =1, [m]], we define the Wasserstein
(IT 1/P HERT
(P(PQQ)=
min
III [wi...., w.)) 11Th [w''.
where D E R is the distance matrix with Di, = dx (xx( One can think of these weighted point sets as discrete probability distributions in (X, dx). When p 1, W1 is also commonly known as the Earth Mover's distance (EMD). Additionally, note that when = oc, W is the same the Hausdorff distance between points in P and Q with non-zero weight. Computing Wasserstein distances amounts to solving a linear programming problem, which takes O(N log N) (where N max m)) time. There have been number of methods for fast approximations of Wasserstein distances, including multi-scale and hierarchical solvers [24). and L₁ embeddings via quadtree algorithms [4]. In particular, entropic regularization of Wasserstein distance (6], also known as the Sinkhorn distance, is often used as the standard Wasserstein distance approximation for learning tasks. Unlike Wasserstein distance, the Sinkhorn approximation is differentiable and can be computed in approximately O(n)) time. The computation time is governed by a regularization parameter c. As approaches zero, the Sinkhorn distance approaches the true Wasserstein distance.
3 Learning functions between point sets
We will first present general framework for approximating SFGI-functions and then show how this framework along with geometric sketches of our input data enables us to define neural networks which can approximate p-Wasserstein distances with complexity independent of input data size.
3.1 A general framework for functions on product spaces
One of the key ingredients our approach is the introduction of what we call a sketch of input data an Euclidean space whose dimension independent of the size of the input data.
Definition 3.1 (Sketch) Let 8 > 0, E N+, and G be a group which acts on X. A (5, a, G)-sketch of metric space (X,dx dx) consists of a G-invariant continuous encoding function + Ra and continuous decoding function _ X such that dx h((),, < d.

---

Now let (Ai, dx, ...... (Km, dxx..) be compact metric spaces. The product space X1 x ... x Xmm is still a metric space equipped with the following natural metric induced from metries of each factor:
xx.xx..xx.... (A.....,AA, (A'...., A'm)) = dx; (A1, 41) +... + dx.. (Am A
Suppose G2 is a group acting on N, for each i [m-] In the following result, instead of SFGI product functions, we first consider the more general case of factor-wise group invariant functions, namely functions f Xr x... x Xm + R such that is uniformly continuous and invariant to the group action G1 X ... x Gm-
Lemma 3.2 Suppose f X1 x... Xm + R is uniformly continuous and invariant to G x... X Gmn- Additionally, assume that for any 8 > 0, (X,, dx,) has a (8, ai, Gi)-sketch where may depend on Then for any > O, there is a continuous G--invarian functions Φ; X, + Ra for all E [ki and continuous function p R 19
x mm - IR such that
[A(A, A2 ................................ Am) - (A)) (A2)) ..... 0n ( (Am))) <E
for any (A1...., Am) E X1 x ... x Xm- Furthermore, if X1 = ...22 _... Xm,, then we can choose 11 = 122 .... Pm.
Note that recent result from 17 shows that continuous factor-wise group invariant function I Xi x ......... .,
can be represented (not approximated) by the form f(V,,..., Um) p(δi ((1)( on (v.) if there exists a topological embedding from ///G; to Euclidean space. The condition that each quotient& X./G; has a topological embedding in fixed dimensional Euclidean space is strong. A topological embedding requires injectivity, while in a sketch, one collapse input objects as long as after decoding, we obtain an approximated object which is close to the input. Our result can be viewed as a relaxation of their result by allowing our space to have an approximate diedddimensional embedding (i.e., our (6, a, G)-sktchh..
We often consider the case where X' = X1 = = Xm i.e. f: x .......' R where G' is group acting on the factor X. Oftentimes, we require the function to not only be invariant to the actions of a group G on each individual X' but also symmelric with respect to the ordering of the input. By this, we mean ...,.... Am) = [(Aa(1):--... Ax(m)) where π is a permutation on [m]] other words, consider the SFGI product function f as introduced in Definition
[.I The extra symmetry
requirement adds more constraints to the form of f. We show that the set of uniformly continuous SFGI product function can be universally approximated by product function with even simpler form than Lemma 3.2 as stated in the theorem below.
Lemma 3.3 Assume the same setup Lemma 3.2 with x= x =
Xm and G = =G Gm- Assume that X has a (8, a, G-ssketch. Additionally, suppose f is symmetric; hence f is a there a continuous G-invariant function 6: X + Ie' and continuous functiionn p: + IR such that
SFGI function. Let = 2am + Then for any
[f(A). Am)
((A.)) <E
Now suppose we want approximate an SFGI product function, f. with neural network. Lemma 3.3 implies that we can approximate ó with any universal G-invariant neural network which embeds original space X' to some Euclidean space R Then the outer architecture P can be any universal architecture (e.g. MLP). Finding universal G-invariant neural network to realize & over a single factor space X' is in general much easier than finding a SFGI neural network, and as we discussed at the end of Section [ already know how achieve this for several settings. We will show how this idea is at work for approximating SFGI functions between point sets in the next subsection.
3.2 Universal neural networks for functions between point sets
We are interested in learning symmetric functions between point sets (i.e. any p- Wasserstein distance) which are factor-wise permutation invariant. this section, we will show that we can find (8, a, G)- sketch for the space of weighted point sets. This allows us to combine Lemma 3.3 with DeepSets to define a set of neural networks which can approximate p-th Wasserstein distances to arbitrary accuracy. Furthermore, the encoding and decoding functions can be approximated with neural networks where their model complexity is independent of input point set size. Thus, the resulting neural network used to approximate Wasserstein distance also has bounded model complexity.
5

---

Set up. Given some metric space (92, dg), let X be the set of weighted point sets with at most N elements. In other words, each S E X' has the form S = {(z, w.) w, E t E. w, 1, [t} and |S| < N. One can also consider X' to be the set of weighted Dirac measures over Q. For simplicity, we also sometimes use S to refer to just the set of points {æ:) contained within it. We will consider the metric over X' to be the p-th Wasserstein distance, W,, We refer to the metric space of weighted point sets over as (x, Wp).
First, we will show that given a ò, there is a (6, a, G)-sketh of X' with respect to Wp. The embedding dimension a depends on the so-called covering number of the metric space (92, da) from which points are sampled. Given compact metric space (92, dg). the covering number (22(r) w.t.. radius r is the minimal number of radius r balls needed cover Q. As a simple example, consider ? = i= A] S R. Given any we can cover X with 2A intervals so (r) I 2 The collection
of the center of set of r-balls that cover Ω an r-net of S2. For compact set S C IR with diameter D, its covering number (r) is a constant depending on D, r and d only.
Theorem 3.4 Set dx to be for 1 OC. Let G be the permutation group. For any ó > 0, let ó0 1 5/22 and a vr)(6)) be the covering number w.r.t. radius óg Then there is a (ő, a, G)- sketch of X with respect to W p. Furthermore, the encoding function - R can be expressed as the following where h: Ω - R" is continuous:
h(S) Σ Mxx)
DES
Proof: Let > 0 and let S E X' be S (1x,, wi) ' W; = 1,2; E 2) and |S| S N. Given 5n = 2 1/02 and a = (d)), we know Ω has a dg,net, C. and we denote the elements of C as (V1---- .o.. In other words, for any x E 2 there is a yi E Cd such that da(ar, v1) < 86-
First, we will define an encoding function p : X + R.. For each yr, we will use a soft indicator function e- --bdg(a,Bao (w)) and set the constant b so that e ebdg( Bao (16) is "sufficiently" small
dr(x, B&g (21)) > dg- More formally, we know that lim e-660 0 so there is 8 E R such that
for all > B, e 660 Ta
Set by to be such that bn > 3. Let (z( = e bbda(a, (v)) for each
i [a]. For given x E 12 we compute h Ω + Ra as
h(x) [Ar (x), ... (x)
Then we define the encoding function h: + Ra as
M((a)) Wi [h(x:))l
h(S) = Σ i=1
Note that [|h(S)] = and h is continuous since Wasserstein distances metrize weak convergence. Additionally, since (x, B&, (v+)) is the distance from to the baal around yi, we are guaranteed to have one j where he (ax)) = 1 so [1(x2;)1 >>> 1.
Now, we define a decoding function g : R + X as g(v) ((v, ) i E [a]]. In order to show that g and h yields a valid (6, a, G--sketch of X, we must show that g h(S) sufficiently close to S under the W, distance. First, we know that
hy(2;) [M(x))
<
wi
d(r.,1))
Let dmaa be the diameter of Ω. For a given Ti, we can partition (Ze),...... ha (2:)) into o7 those where hz(xx)) Theas a and those where hy(2;) &h
Tha i.e. (hy (xx)).... hi, (xx) and
(hre-, (2.). ... h2g (x,)) respectively. If hz) >
T Ther then
82
%
badat(a, (v:))
>
6yóo
d a

---

so do (5, (B:)) < 80. Then
(x(()))
W Wo (S)) v
[1((m))]1
(1r))
8
Wi
(260)
pmaz))
d
since dg (2i, Uj) S dmax
[|((z;)]11
a
2=k+1
wi (2" %g + %) 2"
( (0/2
2
+ 2
1
8
+
8
2
2
Thus, the encoding function h and the decoding function make up (6, a, G)-skttch.
Note that the sketch outlined in Theorem 3.4 is smooth version of one-hot encoding, With Theorem 3.4 and Lemma 3.3 we will now give an explicit formulation of an e-approximation of f via sum-pooling of continuous functions.
Let O, (St, da) be a compact metric space and let X be the space of weighted point equipped with the p- Wasserstein, W Suppose for any 5, (S2, da) has covering number a(6). Then given a function f X x X + R that is uniformly continuous and permutation invariant, there is continuous functionss h : Ω + (a()), Ra(8) + R.. and p R. - E, such that for any
Corollary 35
A, B E X
S(A.B) - J(A.B)-
plo(
Σ wr((x)) + Φb
Σ wrb(s))) <E
(z,00)eA
(x,
where h, and p are all continuous and a' 4 a(0) a(8)+1 +1.
Due Eqn. ( instead of considering the function which takes a set of points S X' input, we now only need to model the function / : S2 - IR=, which takes a single point æ E S as input. For simplicity, assume that the input metric space (S2, do) compact set in some Euclidean space IR Note that in contrast to Lemma 3.3 each h, 6 and p simply a continuous function, and there is no further group invariance requirement. Furthermore, all the dimensions of the domain and range of these functions are bounded values that depend only on the covering number of Ω the target additive error e, and independent (MIPs) ha, - Φ0,; and poş in place of h, & and to approximate the desired function. Formally, we define the following family of neural networks:
the maximum size N of input points. We can use multilayer perceptrons
(A, NeroductNer (A, B) = pas ( (doe
Σ (2,0)))A
Wyhe, + +0z2
Σ (B.W))EB
Wyhe, (x)
(2)
In practice, we consider the input the the neural network he, be a point x Ω along with its weight wg. As per the discussions above, functions represented by N NPProductNet can universally approximate SFGI product functions on the space point sets. See Figure f an illustration of our universal architecture for approximating product functions on point sets. As p- Wasserstein distances, Wp X' + R. are uniformly continuous with respect to the underlying metric W, ", we can apply our framework for the problem of approximating p- Wasserstein.
Importantly, the number of parameters in NProduct does not depend on the maximum size of the point set but rather only on the e additive error and by extension, the covering number of the original metric space. This is because the encoding function for our sketch defined as the summation of single points into R where is independent of the size of the input set. Contrast this result with latent dimension in the universality statement of DeepSets (cf. Theorem 2.I. which is dependent on the input point set size. Note that in general, the model complexities of the MLPs Pa,, Φ0,, and he, depend on the dimensions of the domain and co-domain of each function they approximate (p. é and h) and the desired approximation error We assume that MLPs are Lipschitz continuous. In our case, dh, operates on the sum of ha, (ar)s for all number of input points A (or in B). In general, the error made in hs may accumulate N times, which causes the precision we must achieve for each individual ho, (x) (compared to h(x)) to be O(e/N), This would have caused the model complexity of he, to depend on N. Fortunately, this is not the case for our encoding
7

---

Euclidean embeddings from sum pooling over elements
Slamese network ser elements
Siamese network on embeddinge
Final num-pooling and MLP
1 = (
he,
h(4) Σ (mx)
$8.
(b,(b(())
p0s
(d(o)(((AA)))) +o,,(((A)))
ho:
b(B)
Σ (n)
Φo
(b,(b(B)
Figure 1: Visually representing neural network which can universally approximate uniformly continuous SFGI product functions over pairs of point sets.
function h. In particular, because our encoding function can be written in the specific form of normalized sum of individual points in the set S; i.e, de, operates on La,,,,ess Wahe (x) with
E. W/ 1. the error accumulated by the normalized sum will be less than the maximum error from any single point h(x) for z E S. Thus, both the error for each MLP and the dimension of the domain and co-domain the input point set N, we get that P0,- Φo, and he, each have model complexity independent of the size of the input point set. In short, because the encoding and decoding functions our sketch can be approximated with neural networks with model complexity independent of input point set size, we are able to achieve low model complexity neural network which can approximate Wasserstein distance arbitrarily well. Note that in general, can be approximated by neural networks which are independent of input size. Finally, this neural network architecture also be applied to approximating Hausdorff distance. More details regarding Hausdorff distance approximation are available in Appendix A We summarize the above results in the following corollary.
each approximated function (p, o, h) do not depend on the size of
not guaranteed to be able to find sketch which
There a network the form of NPrrddusNett which can approximate the p- Wasserstein distance between point sets within additive c-error. The number of parameters in this network depends on e and the covering number of Ωl and not on the size of each point set.
Corollary 3.6
Additionally, if we replace the sum-pooling with max in NProductNet there is network of such form where we can also approximate Hausdorff distance between point sets to additive e aaccuray.
Exponential dependence on dimension Although the model complexity of NProductNot is inde- pendent of the size of the input point set, it depends on the covering number of Ω which, in turn, have exponential dependence
the dimension of Ω. In short, this means that the model complexity of NPrrodutNet has an exponential dependence on the dimension of Ω. However, in practice, many machine learning tasks (c.g. 3D point processing) involve large point sets sampled from low-dim wwwddimensionll space (d = 3). Furthermore, in general, the covering number of Ω will depend on the intrinsic dimension l rather than the ambient dimension. For instance, if the input point sets are sampled from hidden manifold of dimension d' (where d' which is much lower than the ambient dimension d), then the covering number would depend only on d' and the curvature bound of the manifold. In many modern machine learning applications, it is often assumed that the data sampled from a hidden space of low dimension (the manifold hypothesis) although the ambient dimension might be very high.
Using max-pooling instead of sum-pooling. Observe that in the final step of combining the Euclidean outputs for two point sets A, B X, E(r,e,ed whho (x) + E(E,,,) Weho, (z), we the of these components (as in DeepSet architecture) E(E.w,)EA Waho, (Ir and Exrma e whoa (z), to ensure the symmetric condition SFGI product functions. One could replace this final sum with final max such in PointNet. However, to show that PointNets are able to universally approximate continuous functions F K - R. where K pax2 is compact, need
8

---

Table Mean relative error between approximations and ground truth Wasserstein distance between point sets. The top row for each dataset shows the approximation quality for point sets with input sizes that were seen at training time; while he bottom row shows the approximation quality for point sets with input sizes that were not seen at training time. Note that Net our model.
Dataset
Input size
NeroductNet
WPCE
NeDeepSeta
Sinkhorn
[100, 300]
0.046 1 0.043
0.341 0.202
0.362
0.241
0.187 T 0.232 0.241 1 0.325
noisy-sphere-3
[300, 500]
0.158 1 0.198
0.356 1 0.286
0.608 + 0,431
[100, 300]
0.015 + 0.014
0.269 + 0.285
0.291 = 0,316
0.137 + 0.122
noisy-sphere-6
[300, 500]
0.049 + 0.054
0.423 + 0.408
0.508 + 0.473
0.198 + 0.181
256
0.097 +: 0.073
0.120 + 0.103
0.123 + 0.092
0.073 + 0.009
uniform
[200, 300]
0.131 + 0.096
1.712 + 0.869
0.917 = 0.869
0.064 + 0.008
[20, 200]
0.084 + 0.077
0.077 + 0.075
0.105 + 0.096
0.101 + 0.032
ModeINet-small
[300, 500]
0.111 + 0.086
0.241 + 0.198
0.261 + 0.245
0.193 + 0.155
2048
0.140 + 0.206
0.159 + 0.141
0.166 + 0.129
0.148 + 0.048
ModeINet-large
[1800, 2000]
0.162 + 0.228
0.392 + 0.378
0.470 + 0.628
0.188 + 0.088
[20, 200]
0.012 + 0.010
0.477 + 0.281
0,482 + 0.291
0.040 + 0.009
RNAseq
[300, 500]
0.292 + 0.041
0.583 0.309
0.575 + 0.302
0.048 + 0.006
Table 2: Training time and number of epochs needed for convergence for best model
Dataset
NProduetNet
WPCE
NeDeepSett
Time
6min
Ihr 46min
9min
noisy-sphere-3
Epochs
20
100
100
Time
12min
4hr 6min
lhr 38min 100
noisy-sphere-6
Epochs
20
100
Time
7min
3hr 36min
1hr 27min
uniform
Epochs
23
100
100
Time
7min
1hr 23min
12 min
ModelNet-small
Epochs
20
100
100
Time
8min 20
3hr 5min 100
40min
ModeINet-large
Epochs
100
Time
15min 73
14hr 26min
3hr 01
RNAseq(2k)
Epochs
100
100
to use a ō-net for K which will also serve as the intermediate dimension for S. As K C [0, Njax in our case (where N is the maximum cardinality for a point set), the intermediate dimension for max-pooling invariant architecture at the end (i.e. PointNet) now depends on the maximum size of input point sets.
4 Experimental results
We evaluate the accuracy of the J-Wasserstein distance approximations of our proposed neural network architecture, .NProductNet- against two different baseline architectures: (1) a Siamese autoencoder known the Wasserstein Point Cloud Embedding network (WPCE) [15] (previously introduced at the end of Section | and is SOTA method of neural approximation of Wasserstein distance) and (2) a Siamese DeepSets, denoted as NNDeeppets, which a single DeepSets model which maps both point sets to a Euclidean space and approximates the 1- W-Wasserstein distance as the & norm between output of each point set. As Siamese networks are widely employed for metric learning, NeDerpSets model is natural baseline for comparison again NProduetNet- We additionally test our neural network approximations against the Sinkhorn distance where the regularization parameter was set to 0.1. For each model, record results for the best model according to hyperparameter

---

search with respect to the parameters of each model. Finally, we use the ModelNet40 [3]] dataset which consists of point clouds in R3 and a gene expression dataset (RNAseq) which consists of 4360 cells each represented by 2000 genes (i.e. 4360 points in R2000) as well as three synthetic datasets: (1) uniform, where point sets are in R2, (2) noisy-sphere-3, where point sets are in R&, (3) noisy-sphere-6, where point sets are in R The RNAseq dataset is publicly available courtesy of the Allen institute [32]. Additional details and experiments approximating the 2- 2-aasserteein distance are available in Appendix D
Approximating Wasserstein distances. Our results comparing 1-Wasserstein distance approxima- tions are summarized in Table I Additionally, see Table Tableßoor for a summary of time needed for training. For most datasets, ProductNNet produces more accurate approximations of Wasserstein distances for both input point sets seen at training time and for input point sets unseen at training time. For the high dimensional RNAseq dataset, our approximation remains accurate in comparison with other methods, including the standard Sinkhorn approximation for input point sets seen at training time. The only exception is ModelNet-smaal, where the NPrrouctNet approximation error is slightly larger than WPCE for input point set sizes using during training (top row for each dataset in T. However, for point sets where the input sizes were not used during training (bottom row for each dataset in Table D NVProductNet showed siginificantly lower enror than all other methods including WPCE. These results along with a more detailed plot in Figure 2 in Appendix indicate that NProddutNett generalizes better than WPCE point sets of input sizes unseen at training time. Also, see Appendix for additional discussion about generalization. Furthermore, one major advantage of NProductNet over WPCE the dramatically reduced time needed for training (cf. Table 2). This substantial difference in training time is due to WPCE's usage of the Sinkhorn reconstruction loss as the (n²²) computation time for the Sinkhorn distance can be prohibitively expensive as input point set sizes grow. Thus, our results indicate that, compared to WPCE, NPProducteet can reduce training time while still achieving comparable or better quality approximations of Wasserstein distance. Using our NProductNet, we can produce high quality approximations of 1- I-Wasserstein distance while avoiding the extra cost associated with using an autoencoder architecture and Sinkhorn regularization. Finally, all models produce much faster approximations than the Sinkhorn distance (see Tables 3 and 6 in Appendix DD. In summary, as compared to WPCE, our model is more accurate approximating both 1- Wasserstein distance, generalizes better to larger input point set sizes, and is more efficient in terms of training time.
Concluding Remarks
Our work presents general neural network framework for approximating SFGI functions which can be combined with geometric sketching ideas into a specific and efficient neural network for approximating p- p-Wasserstein distances. We intend to utilize WProductNet an accurate, efficient, and differentiable approximation for Wasserstein distance in downstream machine learning tasks where Wasserstein distance is employed, such as loss functions for aligning single cell multi-omics data [8] or compressing energy profiles in high energy particle colliders 10, 26). Beyond Wasserstein distance, we will look to apply our framework to wide array of geometric problems that can be considered SFGI functions and are desireable to approximate via neural networks. For instance, consider the problems of computing the optimal Wasserstein distance under rigid transformation or the Wassoosstnsnrnnnnn distance, which both can be represented as an SFGI function where the factor-wist group invariances include both permutation and rigid transformations. Then our sketch must be invariant to both permutations and orthogonal group actions on the left. It remains to be seen if there is a neural network architecture which can approximate such an SFGI function to within an arbitrary additive c-error where the complexity does not depend on the maximum size of the input set.
6
Acknowledgements
The authors thank anonymous reviewers for their helpful comments. Furthermore, the authors thank Rohan Gala and the Allen Institute for generously providing the RNAseq dataset. Finally, Samantha Chen would like to thank Puoya Tabaghi for many helpful discussions about permutation invariant neural networks and Tristan Brugère for his implementation of the Sinkhorn algorithm. This research supported by National Science Foundation (NSF) grants CCF-2112665 and CCF-2217058 and National Institutes of Health (NIH) grant RFI MH125317.
10

---

References
[1] T. Amir, S. J. Gortler, I. Avni, R. Ravina, and N. Dym. Neural injective functions for multisets, measures and graphs via finite witness theorem. arXiv preprint vrXiv:2306.06522) 2023.
[2] A. Backurs, Y. Dong, P. Indyk, I. Razenshteyn, and T. Wagner. Scalable nearest neighbor search for optimal transport. In International Conference machine learning, pages 497-506. PMLR, 2020.
[3] P. Bongini, M. Bianchini, and F. Scarselli. Molecular generative graph neural networks for drug discovery. Neurocomputing, 450:242-252, 2021.
[4] C. Bueno and A. Hylton. On the representation power of set pooling networks. Advances in Neural Information Processing Systems, 34:17170-17182, 2021.
[5] N. Courty, R. Flamary, and M. Ducoffe. Learning wasserstein embeddings. arXiv preprint 1rXiv:1710.0774 2017.
[6] M. Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. Advances in neural information processing systems, 26, 2013.
[7] G. Cybenko. Approximation by superpositions of a sigmoidal function. Mathematics of control, signals and systems, 2(4):303-314, 1989.
[8] Demetci, R. Santorella, B. Sandstede, W. Noble, and R. Singh. Cromov-wasserstern optimal transport align single-cell multi-omics data. BioRxiv, pages 2020-04, 2020.
[9] C. Deng, O. Litany, Y. Duan, A. Poulenard, A. Tagliasucchi, and L. Guibas. Vector neu- A general framework for (3)-qquivrriant networks. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 12200-12209, 2021.
[10] G. Di Guglielmo, F. Fahim, C. Herwig, M. B. Valentin, J. Duarte, C. Gingu, P. Harris, J. Hirschauer, M. Kwok, V. Lonear, al. reconfigurable neural network asic for detector front- end data compression at the hl-lhc. IEEE Transactions Nuclear Science, 6818):2179-2188 2021.
[11] W. Du, H. Zhang, Y. Du, Q Meng, W. Chen, N. Zheng, B. Shao, and T.-Y. Liu. Se (3) equivariant graph neural networks with complete local frames. International Conference on Machine Learning, pages 5583-5608. PMLR, 2022.
[12] N. Dym and Gortler. Low dimensional invariant embeddings for universal geometric learning. arXiv preprint arXiv:2205.02956 2022.
[13] Gilmer, S. Schoenholz, P. F. Riley, O. Vinyals, and G. E. Dahl. Neural message passing for quantum chemistry. In International conference on machine learning, pages 1263-1272 PMLR, 2017.
[14] Indyk and N. Thaper. Fast image retrieval via embeddings. 3rd international workshop on statistical and computational theories of vision, volume 2, page 5, 2003.
[15] K. Kawano, S. Koide, and T. Kutsuna. Learning wasserstein isometric embedding for point clouds. In 2020 International Conference 3D) Vision (3DV), pages 482..82. IEEE, 2020.
[16] N. Keriven and G. Peyré. Universal invariant and equivariant graph neural networks. Advances in Neural Information Processing Systems, 32, 2019.
[17] D. Lim, J. Robinson, L. Zhao, Smidt, Sra, H. Maron, and S. Jegelka. Sign and basis invariant networks for spectral graph representation learning, arXiv preprint arXiv:2202.13013 2022.
[18] A. Makkuva, A. Taghvaei, S. Oh, and J. Lee. Optimal transport mapping via input convex neural networks. In International Conference on Machine Learning, pages 6672-6681. PMLR, 2020.
11

---

[19] H. Maron, H. Ben-Hamu, N. Shamir, and Y. Lipman. Invariant and equivariant graph networks. arXiv preprint arXiv:1812.09900 2018.
[20] H. Maron, E. Fetaya, N. Segol, and Y. Lipman. On the universality of invariant networks. In International conference machine learning, pages 4363-4371. PMLR, 2019.
[21] H. Maron, O. Litany, G. Chechik, and E. Fetaya. On learning sets of symmetric elements. In International conference on machine learning, pages 7744. PMLR, 2020.
[22] C. R. Qi, H. Su, K. Mo, and I. J. Guibas. Pointnet: Deep learning on point sets for 3d classification and segmentation. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 652-660, 2017.
[23] F. Scarselli, M. Gori, A. C. Tsoi, M. Hagenbuchner, and G. Monfardini. The graph neural network model. IEEE transactions on neural networks, 20 (0(1):61-800) 2008.
[24] B. Schmitzer. A sparse multiscale algorithm for dense optimal transport. Journal of Mathemati- cal Imaging and Vision, 56(2):238 259, 2016.
[25] N. Segol and Y. Lipman. On universal equivariant set networks. arXiv preprint arXiv:1910.00421,,000 2019.
[26] R. Shenoy and The CMS Collaboration Team. EMD Neural Network for HGCAL Data Compression Encoder ASIC. In APS April Meeting Abstracts, volume 2022 of APS Meeting Abstracts, page Q09.008, Jan. 2022.
[27] A. Taghvaei and A. Jalali. 2-wasserstein approximation via restricted convex potentials with application to improved training for gans. arXiv preprint 192.07999 2019.
[28] M. A. Uy, Q.- Pham, B.-S. Hua, T. Nguyen, and S.-K. Yeung. Revisiting point cloud classification: A new benehmark dataset and classification model on real-world data. In Proceedings of the IEEE/CVF international conference on computer vision, pages 1588-1597, 2019.
[29] E. Wagstaff, F. Fuchs, M. Engelcke, Posner, and M. A. Osboree On the limitations of representing functions on sets. In International Conference on Machine Learning, pages 6487-6494. PMLR, 2019.
[30] E. Wagstaff, F. Fuchs, M. Engelcke, M. Osborn,, and Posner. Universal approximation of functions on sets. Journal of Machine Learning Research, 23(151):1-56, 2022.
[31] Z. Wu, S. Song, A. Khosla, F. Yu, L. Zhang, X. Tang, and J. Xiao. 3d shapenets: A deep representation for volumetric shapes. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 1912-1928 2015.
[32] Z. Yao, C. T. van Velthoven, T. N. Nguyen, Goldy, A. E. Sedeno-Cortes, F. Baftizadeh, D. Bertagnolli, T. Casper, M. Chiang, K. Crichton, et al. A taxonomy of transcriptomic cell types across the isocortex and hippocampal formation. Cell, 184(12):3222 321232222,,,,1,,,,,, 2021.
[33] D. Yarotsky. Universal approximations of invariant maps by neural networks. Constructive Approximation, 55(1):407-474, 2022.
[34] M. Zaheer, S. Kottur, S. Ravanbakhsh, B. Poczos, R. R. Salakhutdinov, and A. J. Smola. Deep sets. Advances neural information processing systems, 30, 2017.
12


Reviews:
[{"review_id": "PoLRhFcoRt", "reviewer": "Reviewer_1U84", "summary": "The paper proposes a neural network architecture that efficiently approximates $p$-Wasserstein distance between point sets. This is achieved by exploiting the networks' capability to approximate functions that are symmetric and invariant to group actions componentwise. The highlight of the model is that its complexity remains independent of the input sample sizes, improving generalization capabilities. Supporting empirical evidence also showcases the time efficiency of the approach compared to SOTA architectures. ", "questions": "1. What is meant by '$1$-median' of a collection of point sets? [Line 39] Also, specify the notation $S$ in Definition 3.1.\n\n2. Shouldn't it be $D_{i,j}=d_{X}(x_{i},x'_{j})^{p}$ instead in Line 153?\n\n3. In dimensionality reduction or embedding problems it is often difficult to find out continuous encoding maps. For example, the very problem of Autoencoders is mostly dedicated to ensuring that $d_{\\mathcal{X}}(g \\circ h(\\cdot), \\cdot) <\\delta$. In that light, the assumption that sketches readily exist seems quite strong (as in Lemma 3.2). The paper itself expresses the concern in Line 287. Perhaps the authors can provide some examples to justify. This seems crucial since in most results regarding embedding onto Euclidean spaces (due to Johnson-Lindenstrauss, Bourgain, etc.) the optimal embedding dimension under a given level of distortion depends on the input sample size (cardinality of the finite input space). \n\n4. The Wasserstein distance operates as a metric on the class of probability distributions defined on $X$ (say, $\\mathcal{P}(X)$). So, is it perhaps more appropriate to define $\\mathcal{X}$ as the set of weighted Dirac deltas corresponding to points from $X$?\n", "limitations": "The paper has no immediate negative societal impact. ", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper is well-written, with ample theoretical discussion regarding the problem and sound proofs. The empirical evidence provided supports the theoretical claims to a fair extent. ", "weaknesses": "The dimension of the range of functions, used to approximate uniformly continuous and permutation invariant maps $f(A, B)$, turn out to be independent of the size of input points, rather depend on the covering number of the input space. Under the assumption of compactness, this becomes straightforward as a finite cover is ensured. As such, the real challenge lies in showing the result under unbounded domains, perhaps with decaying tail conditions (e.g. sub-Exponential) on input distributions.  ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your comments!\n\nRegarding your main comment that \"Under the assumption of compactness, this becomes straightforward as a finite cover is ensured.\", we would like to clarify the following: \nNote that there are **two spaces** involved in our problem. The first space $\\Omega \\subset (X, d_X)$ is the metric space from which input point sets are sampled -- in particular: we take $X = \\mathbb{R}^d$, metric $\\mathrm{d}_X$ to be the standard Euclidean distance $\\| \\cdot \\|_2$, and $\\Omega$ to be a compact set in $\\mathbb{R}^d$. \n\nThe second space, denoted by $(\\mathcal{X}, d_{\\mathcal{X}})$ is the space of weighted point-sets (of maximum cardinality $N$, where points are from $\\Omega \\subset (X, d_X)$). This space, $(\\mathcal{X}, d_{\\mathcal{X}})$, is a factor of the metric space where our product function $f(A, B)$ is defined on; that is, $f: \\mathcal{X} \\times \\mathcal{X} \\to \\mathbb{R}$. In our setting, if input points are from $\\Omega \\subset \\mathbb{R}^d$, then $\\mathcal{X} = \\Omega^N \\subset \\mathbb{R}^{N d}$, and note that the metric $d_{\\mathcal{X}} = W_p$, the $p$-th Wasserstein distance for point-sets. \n\nWhile $\\mathcal{X}$ is compact when $\\Omega$ is compact, note that its dimension is $Nd$, which depends on the maximum size of input point-set. We can not directly use a bound of covering number of this space $\\mathcal{X}$ (which will naively depend on the dimension of $\\mathcal{X}$ which is $Nd$) to define a model as we want the number of parameters to be *independent* of the maximum input point set size $N$.Instead, our Theorem 3.4 shows that when the metric for $\\mathcal{X}$ is $d_{\\mathcal{X}} = W_p$ (the $p$-th Wasserstein distance), we can find a sketch whose latent dimension depends on the covering number of $\\Omega$. \nOur notations of $(X, d_X)$ and $(\\mathcal{X}, d_{\\mathcal{X}})$ might have caused confusion. We will update them to $(\\Omega, \\ell_2)$ and $(\\mathcal{X}, W_p)$ in the revised manuscript to make the distinction more clear.\n\nFurthermore, the existence of a low-dimensional latent space **does not** necessarily imply that one can approximate the map from the original domain to this latent space via a model of bounded complexity. \nTo this end, note that our proof in Theorem 3.4 is constructive: it has specific form using a bounded number of functions, $h_i$'s (for $i\\in [1, a]$), each of which has a simple closed form and can be replaced by a MLP of complexity independent to input point set size.  \n\nRegarding \"What is meant by '1-median' of a collection of point sets?\":\nGiven a collection of $k$ objects $\\{ Y_1, \\ldots, Y_k\\}$ sampled from a metric space $(\\mathcal{Y}, d_{\\mathcal{Y}})$, its {\\it 1-median}, or geometric median, is defined as\n$$ Y^* = argmin_{Y} \\sum_{i=1}^k d (Y, Y_k),$$\nwhich intuitively, is the center of these $k$ objects. For example, if each $Y_i$ is a point set, one can think of $Y^*$ as its average point-set. In our sentence (line 39), we in fact refer to the total distance to the 1-median (or one can think of this as the cost of 1-median), namely, \n$$ cost(Y_1, \\ldots, Y_k) = \\min_Y \\sum_{i=1}^k d (Y, Y_k) = \\sum_{i=1}^k d (Y^*, Y_k).$$\nIf we use squared distance $d^2 (Y, Y_k)$ in the above definitions, then we obtain 1-mean (sometimes also called the Fréchet mean) and its corresponding cost is Fréchet variance. Both 1-median (and its associated cost) and the 1-mean (and the associated Fréchet variance) are commonly used in statistical analysis of collections of objects $Y_i$s (e.g., computing mean / variance in the space of pointsets). While we didn't explicitly state it in the paper, our results / architectures also hold for computing the cost of 1-median or Fréchet mean of a fixed set of pointsets sampled from a compact set in $\\mathbb{R}^d$. In general, there is no closed form solution for 1-median or 1-mean for complex spaces, and they are often computed empirically by solving a non-convex optimization procedure. \n\nRegarding your comment \"In dimensionality reduction or embedding problems ...\": We note that having a sketch for space of point-sets (under Wasserstein distance or some other distance) is different from standard dimensionality reduction or embedding, which often aims to preserves distance metric in the latent space. In our case, we only need the latent space to provide sufficient information so that the Wasserstein distance can be approximately **after decoding**. There is no distance preservation in the latent space.\nAt least for the case of Wasserstein distance, it turns out that it is sufficient to achieve this by mapping to $\\mathbb{R}^p$ where $p$ depends only on the covering number of the input domain. A low covering number (and by extension, a lower dimension for the latent space $\\mathbb{R}^p$) can be obtained if the domain $\\Omega$ has low ``intrinsic dimension'', such as a low-dimensional Euclidean space $\\mathbb{R}^d$ in the linear case, or a hidden non-linear manifold with low intrinsic dimension $d$ and bounded curvature (and the ambient dimension could be high), or a collection of constant number of manifold pieces. \nIn our paper, we use compact subsets of Euclidean space $\\mathbb{R}^d$ for simplicity, but our results can be extended to points sampled from a hidden manifold of dimension $m$ with bounded curvature -- in which case the covering number is also bounded (e.g., using results of [Roër 2013]).  \n\nWe note that the existence of a low-dimensional hidden structure in data is an important assumption for modern machine learning in many settings. For instance, while the space of natural images has very high dimension (size of an image), the intrinsic dimension of the hidden space where real images lie is far smaller. \n\nThank you for your suggestion of defining $\\mathcal{X}$ as weighted Dirac measures. We will incorporate that in the revision."}, {"review_id": "JfUWwxvkRA", "reviewer": "Reviewer_CxcX", "summary": "The goal of this submission is to learn the (Wasserstein) distance between complex objects (e.g. point sets, graphs) within an arbitrary additive $\\epsilon$-error.\n\nThis paper presents a general neural network architecture for approximating symmetric and factor-wise group invariant (SFGI) functions.\n\nThe proposed SFGI neural network could achieve universality with the number of parameters independent of the input point set size.\n\nThe experiment results show the proposed SFGI neural network could approximate the (best or comparable) Wasserstein distance between point sets.", "questions": "in table 1, the row of 'uniform', it seems Sinkhorn achieved a better Wasserstein distance estimation rather then the proposed model, so the wrong highlight.\n\nAny insights for different distances and different complex objects (e.g. graph)? Should the proposed model to change accordingly?", "limitations": "NA", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "This paper presents a general neural network framework for approximating the Wasserstein distance between point sets. \n\nThe size of the proposed universal neural network approximator is independent w.r.t. the input size, which is new and important. The above advantage is demonstrated in Table 2 in the main paper.\n\nThe proposed neural network approximator is simple but may inspire dozens of downstream network designs.\n\nThe presentation is clear. Well written.\n\nThe code is provided.\n\nI haven't checked the supplementary carefully.", "weaknesses": "The proposed neural uses sum-pooling, so that the parameter size will be independent of the input point set size. I am looking forward to seeing more detailed discussions of it. e.g. How could the proposed neural deal with varying input point set size (which is common in many applications)?\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "\nThank you for your comments! \n\nRegarding your question on the handling of **points of varying sizes** (from your comment ``I am looking forward to seeing more detailed discussions of it ...''): First, in our theoretical results, the fact that the parameter size (model complexity) is independent of input point size essentially is from the constructive proof of Theorem 3.4. The proof shows not only that a latent space of dimension independent to $n$ exists, but also that this can be achieved by a bounded set of functions $h_i$s (forming the encoder), for $i\\in [1, a]$, each of which has a simple form and will later be replaced by a small MLP of constant size in our final neural model $\\mathcal{N}_{ProductNet}$. \nNote that a normalization is happening as the input is a weighted point set with total weight $1$ -- one can think that if there are $n$ input points in a point-set, each point receives weight $1/n$. \n\nWhile we assume a maximum size of input pointsets for simplicity in formulating and stating our results, our final neural model can handle any input pointset size (similar to DeepSet, PointNet, or message passing GNNs) -- indeed, in our experiments, we often train on point-sets of small sizes (100 to 300 points) and show that the final model generalizes very well to point sets of much larger sizes unseen during the training (see the second row for each dataset in Table 1): we tested up to around 2000 points -- as the accuracy degrades only very slightly with larger point sets, we expect that it perform similarly for much larger sizes as well. \n\nWe will make these points more clear in our revised manuscript.\n\nRegarding ''Any insights for different distances and different complex objects (e.g., graphs) ...'': That is a very good question. In the appendix, we showed that the same framework can work for Hausdorff distance (but with max pooling, similar to PointNet). In general, we think that we might be able to extend these to geometric matching such as Fréchet distance for curves (and variance of Fréchet distance) without much change of the architecture. A more challenging next step is to compute optimal Wasserstein distance between point-sets under rigid transformations: While there are ways to handle rotation invariance (basis invariance, e.g., via the approach of [Lim et al., ICLR 2023]), how to do so with bounded model complexity (for the encoder) is still challenging; but we think this is achievable. \nDeveloping such a neural estimation for graph distances is also very interesting -- with suitable distances for the space of graphs and proper assumptions, one might be able to show the existence of low dimensional sketch. However, arguing that the encoder needed has bounded model complexity is more challenging. \n\nFinally, while it is not stated in the paper, our current framework can be directly extended to compute the Fréchet variance or the average distance to 1-median of a set of $k$ pointsets $P_1, \\ldots, P_k$ (see the response to Reviewer 1U84 for the definition of Fréchet mean/Fréchet variance, or 1-median and average distance to it). Roughly speaking, the Fréchet mean of a set of point-sets can be thought of as their ''center'' minimizing the Fréchet variance (which is the total squared distance from the mean-pointset to each input pointset). \n\nWe will add more discussions in the revised manuscript. "}, {"review_id": "stdVq47g6M", "reviewer": "Reviewer_xMX9", "summary": "The paper introduces the concept of symmetric and factor-wise group invariant functions (SFGI), which are continuous and symmetric product functions on complex objects like point sets and graphs. The authors propose a general neural network architecture for approximating SFGI functions and combine it with a sketching idea to develop a specific and efficient neural network that approximates the p-th Wasserstein distance between point sets. The key contribution is that the model complexity of the proposed neural network is independent of the input point set sizes, which is a novel result in the field. The paper demonstrates the effectiveness of the neural network architecture through empirical evaluations, comparing it with other models including a state-of-the-art Siamese Autoencoder. The proposed architecture outperforms existing models in terms of generalization and training speed. The authors highlight the potential application of their framework in solving various geometric optimization problems. ", "questions": "See Section 'Weaknesses'.", "limitations": "The authors should devote more space to the limitations of the proposed approach.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. *Novel Contribution*: The paper introduces the concept of symmetric and factor-wise group invariant functions (SFGI) and presents a general neural network architecture for approximating these functions. This contribution addresses the need for distance functions on complex objects that are invariant to various group actions, such as permutation or rigid transformation.\n\n2. *Bounded Model Complexity*: The authors demonstrate that their proposed neural network architecture for approximating SFGI functions has a bounded model complexity. This result is significant because it shows that there exists a neural network with the capacity to approximate the p-th Wasserstein distance without the complexity depending on the sizes of the input point sets. \n\n3. *Integration of Sketching Ideas*: The paper integrates sketching ideas with the general neural network architecture to develop an efficient and specific neural network for approximating the p-Wasserstein distance between point sets. This integration adds a practical element to the theoretical framework and contributes to the overall efficiency of the approach.\n\n4. *Potential Applications*: The authors discuss potential applications of their framework beyond Wasserstein distance estimation. They highlight its potential use in solving a broad range of geometric optimization problems, such as k-means in a metric space. \n", "weaknesses": "1. *Limited Comparison*: While the paper compares the proposed neural network architecture with a state-of-the-art Siamese Autoencoder and other models, the comparison may not be comprehensive enough. Authors consider only the evaluation of the accuracy of the 1-Wasserstein distance. It would be beneficial to include also case p=2  to better understand the relative strengths and weaknesses of the proposed approach.\n\n2. *Limited Generalizability*: While the proposed framework shows promising results for approximating the p-Wasserstein distance between point sets, the paper does not extensively explore its generalizability to other types of complex objects or distance functions. Further investigation into the applicability of the framework to a broader range of geometric matching problems would strengthen the overall contribution.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "\nThank you for your comments!\n\nRegarding your comments on **limited comparison**: Thank you, and we have now also carried out experiments w.r.t. 2-Wasserstein distance. See our results in the attached PDF. As we reported at the beginning in the general comments to all reviewers, we can see that the improvements over all other ML methods in terms of both accuracy and speed are very similar to the 1-Wasserstein distance. Note that for the Sinkhorn distance, there is a tradeoff of time complexity versus accuracy via the regularization parameter $\\epsilon$. We choose a sufficiently small $\\epsilon$ so that Sinkhorn has high accurary as shown in Table 1 ($\\epsilon = 0.01$). However, as we show in Table 2, the speed is much more slower than ours. In fact, if we increase the regularization parameter $\\epsilon$ to $0.10$, Sinkhorn is still slower than ours even on small dataset, yet the accuracy is already much worse than our neural approximation $\\mathcal{N}_{\\mathrm{ProductNet}}$. On average, the Sinkhorn with respect to $\\epsilon = 0.01$ is 20 times slower than our method on datasets with small input point set size (100-300) and 80 times slower than our method on datasets with larger input point set size (400-600). Note that the gap is increasing as the Sinkhorn distance takes quadratic time to compute. \n\nTo further address your point regarding limited comparison, we also show the effectiveness of our method for higher dimensional data. Here we take a single cell gene expression data set $P$ consisting of $4441$ cell each represented by $4000$ genes (i.e., $4441$ points in $\\mathbb{R}^{4000}$ space. We subsample 3000 pairs of point sets of size 20 to 200 for training, and the testing set consists of 300 pairs of point sets of size 20 to 200. As we can see, our neural estimation remains accurate (see the comparison with Sinkhorn). Unfortunately, due to constraints on computational resources and time, we were not able to generate comparisons to other ML methods (WPCE, $\\mathcal{N}_{\\mathrm{SDeepSets}}$) for RNAseq. \nSee the attached PDF for the additional experimental results, and we will add these results to the revised manuscript.\n\nRegarding your comments on **limited generalizability**: First, we want to point out that just having an efficient and effective neural model to estimate the Wasserstein distance itself is of great interests, which have already attracted various past work (in addition to what we have in Related work of our manuscript, there are also several very recent pieces of work emerging). \n\nOur response to Reviewer CxcX is also relevant, which for completeness we include here as well: \n\nIn the appendix, we showed that the same framework can work for Hausdorff distance (but with max pooling, similar to PointNet). In general, we think that we might be able to extend these to geometric matching such as Frechet distance for curves (and variance of Frechet distance) without much change of the architecture. A more challenging next step is to compute optimal Wasserstein distance between point-sets under rigid transformations: While there are ways to handle rotation invariance (basis invariance, e.g., via the approach of [Lim et al., ICLR 2023]), how to do so with bounded model complexity (for the encoder) is still challenging; but we think this is achievable. \nDeveloping such a neural estimation for graph distances is also very interesting -- with suitable distances for the space of graphs and proper assumptions, one might be able to show the existence of low dimensional sketch, but how to argue that the encoder needed has bounded model complexity is more challenging. \n\nFinally, while it is not stated in the paper, our current framework can be directly extended to compute the Frechet variance, or, the average distance to 1-median of a set of $k$ pointsets $P_1, \\ldots, P_k$ (see the response to Reviewer 1U84 for the definition of Frechet mean/Frechet variance, or 1-median and average ditsance to it). Roughly speaking, the Frechet mean of a set of point-sets can be thought of as their ``center'' minimizing the Frechet variance (which is the total squared distance from the mean-pointset to each input pointset). \n\nWe will add more discussions in the revised manuscript. "}, {"review_id": "8bnjWn1KmA", "reviewer": "Reviewer_vLSh", "summary": "The authors propose a framework for learning Wasserstien distances and other `SFGI functions'. Two key ingredients in their approach are the characterization as all SFGI functions (a universality theorem), and a sketching mechanism which shows that the size of certain components in the networks need not depend on the number of points in a set, but only on the dimension of the points and the required approximation error. Empirical verification of the value of this approach is presented.\n", "questions": "Suggestions:\nExplain a bit more re why to learn Wasserstien distances. Mention you results in the appendix which show better inference time than Sinkhorn.\n\nIn Theorem 2.1: the output dimension of phi has considerably been improved recently. The dimension can be 2kn+1 rather than (n+k) choose k. See Corollary 2.2 in [1] and Theorem 3.3 in [2] (full refs below). Presumably this can lead to improved complexity in the subsequent theorem you prove which rely on Theorem 2.1. \n\nLine 127: not sure that the infinity norm is well defined, unless you specify that F_1 and F_2 are subsets of the same C(X,R).\n\nDefinition of W_p: D should be replaced by D^p. Also the index of w_n' should be w_m'.\n\nDefinition 3.1: specify the group G acts on X.\n\nLemma 3.2 and elsewhere: do you need to specify that f is uniformly continuous (rather than just continuous) if we are working on compact metric spaces?\n\nLine 185: `then we have that'--> then we can choose\n\nLine 191: not clear why the assumption of a topological embedding is strong, given Theorem 2.1. Is it because injectivity is not the same as a topological embedding? Or maybe the issue is the dimensionality?\n\nLine 208: i.e., MLP should be e.g. MLP\n\nLine 223 and Theorem 3.4: I would consider just saying W_p is the distance without introducing d_{fancy X} which is confusing since there is also d_X\n\nI would view the construction in Theorem 3.4 as a smoothened version of one-hot encoding. You might want to mention this connection. \n\nLine 238: M was N previously, better to be consistent.\n\nLine 244: I'm not sure whether d_max was defined? (presumably it is the diameter of X)\n\nThe first inequality just under 254 could be explained more: you're basically saying that some matrix is a transport plan between the measures defined by S and g(h(S)).\n\nThe second line in said inequality has an explanation in the middle, this could be styled better.\n\nIn Corollary 3.5: you do not specify that delta depends on epsilon. The quantifiers should be reoredered. Also note that for the f you are interested in, namely the W_p distance, f is actually 1-Lipschitz so you can take delta=epsilon.\n\nCorollary 3.6: mention dependence on dimension not only on epsilon.\n\nLine 354: wasn't completely clear whether this was referring to the permutation+orthogonal problem from the previous line or this was a new problem which I didn't fully understand\n\n \n\n[1] Low Dimensional Invariant Embeddings for Universal Geometric\nLearning, https://arxiv.org/pdf/2205.02956.pdf\n[2] Neural Injective Functions for Multisets, Measures\nand Graphs via a Finite Witness Theorem https://arxiv.org/pdf/2306.06529.pdf", "limitations": "Exponential dependence on dimension of points in the set should be mentioned.", "rating": 6, "confidence": 5, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "The paper presents two nice ideas:\n(a) Characterizaton of SFGI functions (not so suprising in my humble opinion, but still new)\n(b) A `sketching' mechanism for the Wasserstein distance which is nice. I think it is worth noting in the paper that the fact that the (psuedo)inverse g of the encoding h can be computed is an advantage over existing encoding which can be proved to be invertible, but we have no good understanding of how to invert them.\n  \nAdditionally the empirical results seem promising ", "weaknesses": " The complexity of the proposed sketching algorithm does not depend on the number of points n, but does depend exponentially on the dimension of the points d. This means that the bounds obtained by the paper would improve upon the bounds obtained by fully injective mappings only for specific values of (d,n,epsilon), namely when epsilon^{-d} is small relative to n. Since in applications often d=3 and n is large, this regime is indeed of interest.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "\nThank you for your very helpful comments! We will accommodate those in our revised manuscript (including mentioning explicitly the dependence on (the intrinsic) dimension). Below we will just mention a couple more major ones. \n\n- Regarding your comment on **exponential dependence on dimension of input points**: Indeed, that is the case. However, note that in general, this is the ''intrinsic dimension\" of the domain where points are sampled from, instead of the ``ambient dimension\". For example, if the input points are sampled from a hidden manifold of dimension $m$, then the covering number will depend only on $m$ and the curvature bound of the manifold. One could also get a bound on a collection of manifold pieces. In other words, it really depends on the intrinsic dimension of the hidden domain points are sampled from, although for simplicity, we use the Euclidean dimension in our paper. We also note that in modern machine learning, we often assume (explicitly or implicitly) that data is sampled from a hidden space of low dimensions (e.g., in the manifold hypothesis), even though the ambient dimension might be very high. We will make this point more clear and explicit in the revised manuscript. \n\n- Thank you for your references regarding improved bounds for (multi)set functions. We will include them into the revised manuscript. Indeed, using results from the provided references, the dimension of the latent space in Lemma 3.3 becomes $2\\cdot a \\cdot m + 1$ and the dimension of the latent space in Corollary 3.4 becomes $2\\cdot (a(\\delta)) \\cdot 2 + 1$. \n\n- Regarding **Line 191**, what we meant is that a topological embedding requires injectivity, while in a sketch, one can collapse input objects as long as after decoding, we obtain an approximated object close to input. "}]

Author rebuttal:
We thank all reviewers for their valuable comments and feedback. 

We are happy that reviewers appreciate our contribution of a simple neural network model (constructed based on our theoretical insights) which can estimate Wasserstein distance accurately with bounded model complexity. 
Indeed, identifying the right neural model for a given problem, especially for an optimization problem, is not often an easy task and the role of representation learning is crucial, especially for complex objects (point-sets in our case). 
At a high level, we are interested in problems relating to geometric matching of complex objects, and we focus on the Wasserstein distance between point-sets in this paper.
Our theoretical analysis implies that a simple neural model consisting of a Siamese architecture **followed by MLP layers at the end (which is crucial)** turns out to be universal with respect to the Wasserstein distance problem and the model complexity is independent of input point sets (see Corollary 3.6 and Figure 1). 
Empirically, we emphasize that our neural model outperforms both Sinkhorn and SOTA ML-based approaches for Wasserstein distance computation in terms of **both accuracy and speed**, and also generalizes significantly better than SOTA ML-based approaches.  

To further demonstrate the effectiveness of our model, as suggested by Reviewer xMX9, we also:

(1) carry out experiments evaluating the accuracy of our approximation for $W_2$, i.e., $p$-th Wasserstein distance for $p = 2$; and 

(2) further tested our accuracy on a high dimensional real RNAseq dataset (in 4000 dimensions), which are gene expression data. 


See the attached pdf for results. 
For the case of 2-Wasserstein distance, Table 1 in the pdf shows a similar trend the 1-Wasserstein distance originally reported in our paper. In fact, the improvement in terms of accuracy over other ML models are even bigger. 
Note that for the Sinkhorn distance, there is a tradeoff of time complexity versus accuracy via the regularization parameter $\epsilon$. We choose a sufficiently small $\epsilon$ so that Sinkhorn has high accurary as shown in Table 1. However, as we show in Table 2, the speed is much more slower than ours. In fact, if we increase the regularization parameter $\epsilon$ to $0.10$, Sinkhorn is still slower than ours even on small dataset, yet the accuracy is already much worse than our neural approximation $\mathcal{N}_{ProductNet}$. On average, the Sinkhorn with respect to $\epsilon = 0.01$ is 20 times slower than our method on datasets with small input point set size (100-300) and 80 times slower than our method on datasets with larger input point set size (400-600). Note that the gap is increasing as the Sinkhorn distance takes quadratic time to compute. 

We also point out that a neural approximation of Wasserstein distance itself, or more generally, optimal transport, is interesting on its own and has already attracted extensive past work (see our related work section), as well as recent new work.
In our Supplementary material, we show that our neural model can be easily extended to Hausdorff distance. Interesting future work will include extending the model to estimate Fréchet distance (or variants) for curves, or optimal Wasserstein / Hausdorff distance under rigid transformations. 

Meta-review: In this paper the authors propose a novel way to approximate Wasserstein distance with neural networks. The reviewers found the paper interesting but had some concerns about the assumptions, the dependence of the method to input size and about the experiments.

The authors did a detailed rebuttal to those concerns with new experiments, and clarifications that was appreciated by reviewers. All reviewers kept their borderline positive score and are OK for accepting the paper that brings a novel architecture for efficient approximation of the Wasserstein distance to the community.

The authors must take into account the reviewers comments and integrate the results and discussions (exponential dependence of the dimension of the space for instance) from the rebuttal in the final version of the paper.
Decision: Accept (poster)