Title: Neural approximation of Wasserstein distance via a universal architecture for symmetric and factorwise group invariant functions
Venue: NeurIPS 2023
Accepted: yes
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

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 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 a 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, 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 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 a broad range of geometric optimization problems (e.g., k -means in a metric space).

1 Introduction

Recently, significant interest in geometric deep learning has led to a focus on neural network architectures which learn functions on complex objects such point clouds [34, 22] and graphs [23]. Advancements in the development of neural networks for complex objects has led to progress in a variety of applications from 3D image segmentation [28] to drug discovery [3, 13]. 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 are 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, 34, 19, 20, 25, 16, 33, 29, 30, 4]

Yusu Wang

However, in many geometric or graph optimization problems, our input goes beyond a single complex object, but multiple complex objects. For example, the p -Wasserstein distance W p ( X,Y ) between two point sets X and Y sampled from some metric space (e.g., R d ) is a function over pairs of point sets. To give another example, the 1 -median of the collection of k point sets P 1 , . . . , P k in R d 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 R d . Then the p -Wasserstein distance can be viewed as a function W p : X × X → R . Similarly, 1-median for k point sets can be modeled by a function from the product of k copies of X to R . Such functions are not only invariant to permutations of the factors of the product space (i.e. W p ( X,Y ) = W p ( Y, X ) ) but are also invariant or equivariant with respect to certain group actions for each factor. For p -Wasserstein distance, W p is invariant to 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 a type of function which we denote as an SFGI product functions.

Definition 1.1 (SFGI product function) Given compact metric spaces ( X i , d X i ) where i ∈ [ k ] , we define symmetric and factor-wise group invariant (SFGI) product functions as f : X 1 ×X 2 ×··· X k → R where f is (1) symmetric over the k factors, and (2) invariant to the group action G 1 × G 2 ×···× G k for some group G i acting on X i , 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 a 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-dimensional 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 to 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 -Wasserstein distance at similar error threshholds. Moreover, our neural network trains much faster than the SOTA Siamese autoencoder. Overall, our 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 2Wasserstein distance[18, 27]. However, for this approach, training is done per pair of inputs and is restricted to the 2-Wasserstein distance which makes it unsuitable for a general neural network approximation of p -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 a compact metric space to within ϵ -accuracy. Siamese networks are another approach for popular approach for learning Wasserstein distances. Typically, a Siamese network is composed of a single neural network which maps two input instances to Euclidean space.

The output of the network is represented then by ℓ p -norm between the output embeddings. In [5], the authors utilize a Siamese autoencoder which takes two histograms (images) as input. For their architecture, a 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 (KL) 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-Wasserstein distance is the ℓ 2 norm between the embeddings. The idea of learning Wasserstein distances via Siamese autoencoders was extended in [15] 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, including [34, 22, 19, 17, 9]. For instance, DeepSets [34] and PointNet [19] 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 [11, 21] or combining basis invariance with permutation invariance [17]. Our work differs from previous work in that we address a 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 as a 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 -Wasserstein distances.

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 a set of elements and a binary operation that satisfies a specific set of axioms: (1) the associative property, (2) the existence of an identity element, and (3) existence of inverses for each element in the set. Given a metric space ( X , d X ) , the action of the group G on X is a function α : G ×X → X that transforms the elements of X for each element π ∈ G . For each element π ∈ G , we will write π · x to denote the action of a group element π on x instead of α ( π, x ) . For example, if G is the permutation group over [ N ] := { 1 , 2 , . . . , N } , and X = R N , then for any π ∈ G , π · x represents the permutation of elements in x ∈ R N via π i.e. given x = ( x 1 , x 2 , . . . , x N ) , π · x = ( x π (1) , x π (2) . . . , x π ( N ) ) . A function f : X → Y is G -invariant if for any X ∈ X and any π ∈ G , we have that f ( X ) = f ( π · X ) .

Universal Approximation. Let C ( X , R ) denote the set of continuous functions from a metric space ( X , d X ) to R . Given two families of functions F 1 and F 2 where F 1 ⊆ F 2 and F 1 , F 2 ⊆ C ( X , R ) , we say that F 1 universally approximates F 2 if for any ϵ > 0 and any f ∈ F 2 , there is a g ∈ F 1 such that ∥ g -f ∥ ∞ < ϵ . Different norms on the space of functions can be used, but we will use L ∞ norm in this paper, which intuitively leads to additive pointwise-error over the domain of these functions. The classic universal approximation theorem for multilayer perceptrons (MLPs) [7] states that a feedforward neural network with a single hidden layer, using certain activation functions, can approximate any continuous function to within an arbitrary additive ϵ -error.

Permutation invariant neural networks for point cloud data One of the most popular permutation 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 a finite multiset S = { x : x ∈ R d } (meaning that an element can appear multiple times, and the number of times an

element occurs in S is called its multiplicity ). The DeepSets model is defined as

where h θ 1 and g θ 2 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 ([34, 1, 12]) Assume the elements are from a compact set in R k , and the input multiset size is fixed as N . Let t = 2 kN + 1 . Then any continuous multiset function, represented as f : R k × N → R which is invariant with respect to permutations of the columns, can be approximated arbitrarily close in the form of ρ ( ∑ x ∈ X ϕ ( x ) ) , for continuous transformations ϕ : R k → R t and ρ : R t → R .

While universality for the case when k = 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 t = ( k + N k ) . It was recently shown in [12, 1] that the output dimension of ϕ can be reduced to 2 kN +1 , which is the dimension of t which we use in Theorem 2.1 and subsequent theorems. In both the cases where the output dimension of ϕ is t = ( k + N k ) or t = 2 kN +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, d X ) be a metric space. For two weighted point sets P = { ( x i , w i ) : x i ∈ X, ∑ w i = 1 , i ∈ [ n ] } and Q = { ( x ′ i , w ′ i ) : x ′ i ∈ X, ∑ w i = 1 , i ∈ [ m ] } , we define the Wasserstein distance between P and Q as

where D ∈ R n × m + is the distance matrix with D i,j = d X ( x i , x ′ j ) . One can think of these weighted point sets as discrete probability distributions in ( X, d X ) . When p = 1 , W 1 is also commonly known as the Earth Mover's distance (EMD). Additionally, note that when p = ∞ , W p is the same as 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 3 log N ) (where N = max { n, m } ) time. There have been a number of methods for fast approximations of Wasserstein distances, including multi-scale and hierarchical solvers [24], and L 1 embeddings via quadtree algorithms [2, 14]. 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 2 ) time. The computation time is governed by a regularization parameter ϵ . As ϵ approaches zero, the Sinkhorn distance approaches the true Wasserstein distance.

3 Learning functions between point sets

We will first present a 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 in our approach is the introduction of what we call a sketch of input data to an Euclidean space whose dimension is independent of the size of the input data.

Definition 3.1 (Sketch) Let δ > 0 , a ∈ N + , and G be a group which acts on X . A ( δ, a, G ) -sketch of a metric space ( X , d X ) consists of a G -invariant continuous encoding function h : X → R a and a continuous decoding function g : R a →X such that d X ( g ◦ h ( S ) , S ) < δ .

Now let ( X 1 , d X 1 ) , . . . , ( X m , d X m ) be compact metric spaces. The product space X 1 ×··· × X m is still a metric space equipped with the following natural metric induced from metrics of each factor:

Suppose G i is a group acting on X i , 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 : X 1 ×··· × X m → R such that f is uniformly continuous and invariant to the group action G 1 ×··· × G m .

Lemma 3.2 Suppose f : X 1 ×···×X m → R is uniformly continuous and invariant to G 1 ×···× G m . Additionally, assume that for any δ > 0 , ( X i , d X i ) has a ( δ, a i , G i ) -sketch where a i may depend on δ . Then for any ϵ > 0 , there is a continuous G i -invariant functions ϕ i : X i → R a i for all i ∈ [ k ] and a continuous function ρ : R a 1 ×··· × R a m → R such that

for any ( A 1 , . . . , A m ) ∈ X 1 ×··· × X m . Furthermore, if X 1 = . . . X 2 = · · · = X m , then we can choose ϕ 1 = ϕ 2 = . . . ϕ m .

Note that a recent result from [17] shows that a continuous factor-wise group invariant function f : X 1 × · · · X m → R can be represented (not approximated) by the form f ( v 1 , . . . , v m ) = ρ ( ϕ 1 ( v 1 ) , . . . , ϕ k ( v m )) if there exists a topological embedding from X i /G i to Euclidean space . The condition that each quotient X i /G i has a topological embedding in fixed dimensional Euclidean space is strong. A topological embedding requires injectivity, while in a sketch, one can 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 fixed-dimensional embedding (i.e., our ( δ, a, G ) -sketch ).

We often consider the case where X = X 1 = · · · = X m i.e. f : X × · · · X → R where G is a 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 symmetric with respect to the ordering of the input. By this, we mean f ( A 1 , . . . , A m ) = f ( A π (1) , . . . , A π ( m ) ) where π is a permutation on [ m ] . In other words, we now consider the SFGI product function f as introduced in Definition 1.1. 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 an even simpler form than Lemma 3.2 as stated in the theorem below.

Lemma 3.3 Assume the same setup as Lemma 3.2 with X = X 1 = · · · = X m and G = G 1 = · · · = G m . Assume that X has a ( δ, a, G ) -sketch. Additionally, suppose f is symmetric; hence f is a SFGI function. Let t = 2 am +1 . Then for any ϵ > 0 , there is a continuous G -invariant function ϕ : X → R t and a continuous function ρ : R t → R such that

Now suppose we want approximate an SFGI product function, f , with a neural network. Lemma 3.3 implies that we can approximate ϕ with any universal G -invariant neural network which embeds our original space X to some Euclidean space R a . Then the outer architecture ρ can be any universal architecture (e.g. MLP). Finding a 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 1, we already know how to 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. In this section, we will show that we can find a ( δ, 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 .

Set up. Given some metric space (Ω , d Ω ) , let X be the set of weighted point sets with at most N elements. In other words, each S ∈ X has the form S = { ( x i , w i ) : w i ∈ R + , ∑ i w i = 1 , x i ∈ Ω } and | S | ≤ N . One can also consider X to be the set of weighted Dirac measures over Ω . For simplicity, we also sometimes use S to refer to just the set of points { x i } contained within it. We will consider the metric over X to be the p -th Wasserstein distance, W p . We refer to the metric space of weighted point sets over Ω as ( X , W p ) .

First, we will show that given a δ , there is a ( δ, a, G ) -sketch of X with respect to W p . The embedding dimension a depends on the so-called covering number of the metric space (Ω , d Ω ) from which points are sampled. Given a compact metric space (Ω , d Ω ) , the covering number ν Ω ( r ) w.r.t. radius r is the minimal number of radius r balls needed to cover Ω . As a simple example, consider Ω = [ -∆ , ∆] ⊆ R . Given any r , we can cover X with 2∆ r intervals so ν Ω ( r ) ≤ 2∆ r . The collection of the center of a set of r -balls that cover Ω an r -net of Ω . For a compact set Ω ⊂ R d with diameter D , its covering number ν Ω ( r ) is a constant depending on D , r and d only.

Theorem 3.4 Set d X to be W p for 1 ≤ p < ∞ . Let G be the permutation group. For any δ > 0 , let δ 0 = 1 2 p √ δ/ 2 and a = ν Ω ( δ 0 ) be the covering number w.r.t. radius δ 0 . Then there is a ( δ, a, G ) -sketch of X with respect to W p . Furthermore, the encoding function h : X → R a can be expressed as the following where h : Ω → R a is continuous:

Proof: Let δ > 0 and let S ∈ X be S = { ( x i , w i ) : ∑ w i = 1 , x i ∈ Ω } and | S | ≤ N . Given δ 0 = 1 2 p √ δ/ 2 and a = ν Ω ( δ 0 ) , we know Ω has a δ 0 -net, C , and we denote the elements of C as { y 1 , . . . , y a } . In other words, for any x ∈ Ω , there is a y i ∈ Cd such that d Ω ( x, y i ) < δ 0 .

First, we will define an encoding function ρ : X → R a . For each y i , we will use a soft indicator function e -b d Ω ( x,B δ 0 ( y i )) and set the constant b so that e -b d Ω ( x,B δ 0 ( y i )) is "sufficiently" small if d Ω ( x, B δ 0 ( y i )) > δ 0 . More formally, we know that lim b →∞ e -bδ 0 = 0 so there is β ∈ R such that for all b > β , e -bδ 0 < δ p 0 d p max · a . Set b 0 to be such that b 0 > β . Let h i ( x ) = e -b 0 d Ω ( x,B δ 0 ( y i )) for each i ∈ [ a ] . For a given x ∈ Ω , we compute h : Ω → R a as

Then we define the encoding function h : X → R a as

Note that ∥ h ( S ) ∥ 1 = 1 and h is continuous since Wasserstein distances metrize weak convergence. Additionally, since d Ω ( x, B δ 0 ( y i )) is the distance from x to the δ 0 -ball around y i , we are guaranteed to have one j where h j ( x i ) = 1 so ∥ h ( x i ) ∥ 1 > 1 .

Now, we define a decoding function g : R a →X as g ( v ) = { ( y i , v i ∥ v ∥ 1 ) : i ∈ [ a ] } . In order to show that g and h yields a valid ( δ, a, G ) -sketch of X , we must show that g ◦ h ( S ) is sufficiently close to S under the W p distance. First, we know that

Let d max be the diameter of Ω . For a given x i , we can partition { h 1 ( x i ) , . . . , h a ( x i ) } into those where h j ( x i ) ≥ δ p 0 d p max · a and those where h j ( x i ) < δ p 0 d p max · a i.e. { h j 1 ( x i ) , . . . , h j k ( x i ) } and { h j k +1 ( x i ) , . . . , h j a ( x i ) } respectively. If h j ( x ) ≥ δ p 0 d p max · a , then

so d ( x, B ( y )) < δ . Then

Thus, the encoding function h and the decoding function g make up a ( δ, a, G ) -sketch.

Note that the sketch outlined in Theorem 3.4 is a smooth version of a one-hot encoding. With Theorem 3.4 and Lemma 3.3, we will now give an explicit formulation of an ϵ -approximation of f via sum-pooling of continuous functions.

Corollary 3.5 Let ϵ > 0 , (Ω , d Ω ) be a compact metric space and let X be the space of weighted point sets equipped with the p -Wasserstein, W p . Suppose for any δ , (Ω , d Ω ) has covering number a ( δ ) . Then given a function f : X × X → R that is uniformly continuous and permutation invariant, there is continuous functions h : Ω → R a ( δ ) , ϕ : R a ( δ ) → R a ′ , and ρ : R a ′ → R , such that for any A,B ∈ X

where h, ϕ and ρ are all continuous and a ′ = 4 · a ( δ ) + 1 .

Due to Eqn. (1), instead of considering the function h which takes a set of points S ∈ X as input, we now only need to model the function h : Ω → R a , which takes a single point x ∈ S as input. For simplicity, assume that the input metric space (Ω , d Ω ) is a compact set in some Euclidean space R d . Note that in contrast to Lemma 3.3, each h , ϕ and ρ is 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 ϵ , and independent to the maximum size N of input points. We can use multilayer perceptrons (MLPs) h θ 1 , ϕ θ 2 , and ρ θ 3 in place of h , ϕ and ρ to approximate the desired function. Formally, we define the following family of neural networks:

In practice, we consider the input to the the neural network h θ 1 to be a point x ∈ Ω along with its weight w x . As per the discussions above, functions represented by N ProductNet can universally approximate SFGI product functions on the space of point sets. See Figure 1 for an illustration of our universal architecture for approximating product functions on point sets. As p -Wasserstein distances, W p : X ×X → R , are uniformly continuous with respect to the underlying metric W p , we can apply our framework for the problem of approximating p -Wasserstein.

Importantly, the number of parameters in N ProductNet does not depend on the maximum size of the point set but rather only on the ϵ additive error and by extension, the covering number of the original metric space. This is because the encoding function for our sketch is defined as the summation of single points into R a where a is independent of the size of the input set. Contrast this result with latent dimension in the universality statement of DeepSets (cf. Theorem 2.1), which is dependent on the input point set size. Note that in general, the model complexities of the MLPs ρ θ 3 , ϕ θ 2 , and h θ 1 depend on the dimensions of the domain and co-domain of each function they approximate ( ρ , ϕ and h ) and the desired approximation error ϵ . We assume that MLPs are Lipschitz continuous. In our case, ϕ θ 2 operates on the sum of h θ 1 ( x ) s for all N number of input points in A (or in B ). In general, the error made in h s may accumulate N times, which causes the precision we must achieve for each individual h θ 1 ( x ) (compared to h ( x ) ) to be Θ( ϵ/N ) . This would have caused the model complexity of h θ 1 to depend on N . Fortunately, this is not the case for our encoding

Γ↽ΓLΓQΓDΓOΓΛΓVΓXΓPΓζΓSΓRΓRΓOΓLΓQΓJΓΛΓDΓQΓGΓΛΓ0Γ◁Γ

Γ6ΓLΓDΓPΓHΓVΓH

ΓQΓHΓWΓZΓRΓUΓNΓΛΓRΓQ

ΓVΓHΓWΓΛΓHΓOΓHΓPΓHΓQΓWΓV

Γ↼ΓXΓFΓOΓLΓGΓHΓDΓQ

ΓVΓXΓPΓΛΓSΓRΓRΓOΓLΓQΓJΓΛΓRΓYΓHΓU

ΓVΓHΓW

ΓΛΓHΓOΓHΓPΓHΓQΓWΓV

Γ6ΓLΓDΓPΓHΓVΓH

ΓQΓHΓWΓZΓRΓUΓNΓΛΓRΓQ

ΓHΓPΓEΓHΓGΓGΓLΓQΓJΓV

function h . In particular, because our encoding function can be written in the specific form of a normalized sum of individual points in the set S ; i.e, ϕ θ 2 operates on ∑ ( x,w x ) ∈ S w x h θ 1 ( x ) with ∑ x w x = 1 , the error accumulated by the normalized sum will be less than the maximum error from any single point h ( x ) for x ∈ S . Thus, as both the error for each MLP and the dimension of the domain and co-domain of each approximated function ( ρ, ϕ, h ) do not depend on the size of the input point set N , we get that ρ θ 3 , ϕ θ 2 and h θ 1 each have model complexity independent of the size of the input point set. In short, because the encoding and decoding functions of our sketch can be approximated with neural networks with model complexity independent of input point set size, we are able to achieve a low model complexity neural network which can approximate Wasserstein distance arbitrarily well. Note that in general, we are not guaranteed to be able to find a sketch which can be approximated by neural networks which are independent of input size. Finally, this neural network architecture can 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.

Corollary 3.6 There is a network in the form of N ProductNet which can approximate the p -Wasserstein distance between point sets to within an additive ϵ -error. The number of parameters in this network depends on ϵ and the covering number of Ω and not on the size of each point set.

Additionally, if we replace the sum-pooling with max in N ProductNet , there is a network of such a form where we can also approximate Hausdorff distance between point sets to additive ϵ accuracy.

Exponential dependence on dimension Although the model complexity of N ProductNet is independent of the size of the input point set, it depends on the covering number of Ω which, in turn, can have an exponential dependence on the dimension of Ω . In short, this means that the model complexity of N ProductNet has an exponential dependence on the dimension of Ω . However, in practice, many machine learning tasks (e.g. 3D point processing) involve large point sets sampled from low-dimensional space ( d = 3) . Furthermore, in general, the covering number of Ω will depend on the intrinsic dimension of Ω rather than the ambient dimension. For instance, if the input point sets are sampled from a 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 is 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 , ∑ ( x,w x ) ∈ A w x h θ 1 ( x ) + ∑ ( x,w x ) ∈ B w x h θ 1 ( x ) , we use the sum of these two components (as in a DeepSet architecture) : ∑ ( x,w x ) ∈ A w x h θ 1 ( x ) and ∑ ( x,w x ) ∈ B w x h θ 1 ( x ) , to ensure the symmetric condition of SFGI product functions. One could replace this final sum with a final max such as in PointNet. However, to show that PointNets are able to universally approximate continuous functions F : K → R where K ⊆ R a × 2 is compact, we need

to use a δ -net for K which will also serve as the intermediate dimension for ϕ . As K ⊆ [0 , N ] a × 2 in our case (where N is the maximum cardinality for a point set), the intermediate dimension for a max-pooling invariant architecture at the end (i.e. PointNet) now depends on the maximum size of input point sets.

4 Experimental results

Weevaluate the accuracy of the 1-Wasserstein distance approximations of our proposed neural network architecture, N ProductNet , against two different baseline architectures: (1) a Siamese autoencoder known as the Wasserstein Point Cloud Embedding network (WPCE) [15] (previously introduced at the end of Section 1 and is a SOTA method of neural approximation of Wasserstein distance) and (2) a Siamese DeepSets, denoted as N SDeepSets , which is a single DeepSets model which maps both point sets to a Euclidean space and approximates the 1-Wasserstein distance as the ℓ 2 norm between output of each point set. As Siamese networks are widely employed for metric learning, N SDeepSets model is a natural baseline for comparison again N ProductNet . We additionally test our neural network approximations against the Sinkhorn distance where the regularization parameter was set to ϵ = 0 . 1 . For each model, we record results for the best model according to a hyperparameter

search with respect to the parameters of each model. Finally, we use the ModelNet40 [31] dataset which consists of point clouds in R 3 and a gene expression dataset (RNAseq) which consists of 4360 cells each represented by 2000 genes (i.e. 4360 points in R 2000 ) as well as three synthetic datasets: (1) uniform, where point sets are in R 2 , (2) noisy-sphere-3, where point sets are in R 3 , (3) noisy-sphere-6, where point sets are in R 6 . The RNAseq dataset is publicly available courtesy of the Allen institute [32]. Additional details and experiments approximating the 2-Wasserstein distance are available in Appendix D.

Approximating Wasserstein distances. Our results comparing 1-Wasserstein distance approximations are summarized in Table 1. Additionally, see Table 3 for a summary of time needed for training. For most datasets, N ProductNet 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-small, where the N ProductNet approximation error is slightly larger than WPCE for input point set sizes using during training (top row for each dataset in Table 1). However, for point sets where the input sizes were not used during training (bottom row for each dataset in Table 1), N ProductNet showed siginificantly lower error than all other methods including WPCE. These results along with a more detailed plot in Figure 2 in Appendix D indicate that N ProductNet generalizes better than WPCE to point sets of input sizes unseen at training time. Also, see Appendix D for additional discussion about generalization. Furthermore, one major advantage of N ProductNet over WPCE is 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 O ( n 2 ) computation time for the Sinkhorn distance can be prohibitively expensive as input point set sizes grow. Thus, our results indicate that, compared to WPCE, N ProductNet can reduce training time while still achieving comparable or better quality approximations of Wasserstein distance. Using our N ProductNet , we can produce high quality approximations of 1-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 D). In summary, as compared to WPCE, our model is more accurate in approximating both 1-Wasserstein distance, generalizes better to larger input point set sizes, and is more efficient in terms of training time.

5 Concluding Remarks

Our work presents a 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 -Wasserstein distances. We intend to utilize N ProductNet as 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 a 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 Gromov-Wasserstein distance, which both can be represented as an SFGI function where the factor-wise 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 ϵ -error where the complexity does not depend on the maximum size of the input set.

References

T. Amir, S. J. Gortler, I. Avni, R. Ravina, and N. Dym. Neural injective functions for multisets, measures and graphs via a finite witness theorem. arXiv preprint arXiv:2306.06529 , 2023.

A. Backurs, Y. Dong, P. Indyk, I. Razenshteyn, and T. Wagner. Scalable nearest neighbor search for optimal transport. In International Conference on machine learning , pages 497-506. PMLR, 2020.

P. Bongini, M. Bianchini, and F. Scarselli. Molecular generative graph neural networks for drug discovery. Neurocomputing , 450:242-252, 2021.

C. Bueno and A. Hylton. On the representation power of set pooling networks. Advances in Neural Information Processing Systems , 34:17170-17182, 2021.

N. Courty, R. Flamary, and M. Ducoffe. Learning wasserstein embeddings. arXiv preprint arXiv:1710.07457 , 2017.

M. Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. Advances in neural information processing systems , 26, 2013.

G. Cybenko. Approximation by superpositions of a sigmoidal function. Mathematics of control, signals and systems , 2(4):303-314, 1989.

P. Demetci, R. Santorella, B. Sandstede, W. S. Noble, and R. Singh. Gromov-wasserstein optimal transport to align single-cell multi-omics data. BioRxiv , pages 2020-04, 2020.

C. Deng, O. Litany, Y. Duan, A. Poulenard, A. Tagliasacchi, and L. J. Guibas. Vector neurons: A general framework for so (3)-equivariant networks. In Proceedings of the IEEE/CVF International Conference on Computer Vision , pages 12200-12209, 2021.

G. Di Guglielmo, F. Fahim, C. Herwig, M. B. Valentin, J. Duarte, C. Gingu, P. Harris, J. Hirschauer, M. Kwok, V. Loncar, et al. A reconfigurable neural network asic for detector frontend data compression at the hl-lhc. IEEE Transactions on Nuclear Science , 68(8):2179-2186, 2021.

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. In International Conference on Machine Learning , pages 5583-5608. PMLR, 2022.

N. Dym and S. J. Gortler. Low dimensional invariant embeddings for universal geometric learning. arXiv preprint arXiv:2205.02956 , 2022.

J. Gilmer, S. 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.

P. Indyk and N. Thaper. Fast image retrieval via embeddings. In 3rd international workshop on statistical and computational theories of vision , volume 2, page 5, 2003.

K. Kawano, S. Koide, and T. Kutsuna. Learning wasserstein isometric embedding for point clouds. In 2020 International Conference on 3D Vision (3DV) , pages 473-482. IEEE, 2020.

N. Keriven and G. Peyré. Universal invariant and equivariant graph neural networks. Advances in Neural Information Processing Systems , 32, 2019.

D. Lim, J. Robinson, L. Zhao, T. Smidt, S. Sra, H. Maron, and S. Jegelka. Sign and basis invariant networks for spectral graph representation learning. arXiv preprint arXiv:2202.13013 , 2022.

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.

H. Maron, H. Ben-Hamu, N. Shamir, and Y. Lipman. Invariant and equivariant graph networks. arXiv preprint arXiv:1812.09902 , 2018.

H. Maron, E. Fetaya, N. Segol, and Y. Lipman. On the universality of invariant networks. In International conference on machine learning , pages 4363-4371. PMLR, 2019.

H. Maron, O. Litany, G. Chechik, and E. Fetaya. On learning sets of symmetric elements. In International conference on machine learning , pages 6734-6744. PMLR, 2020.

C. R. Qi, H. Su, K. Mo, and L. 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.

F. Scarselli, M. Gori, A. C. Tsoi, M. Hagenbuchner, and G. Monfardini. The graph neural network model. IEEE transactions on neural networks , 20(1):61-80, 2008.

B. Schmitzer. A sparse multiscale algorithm for dense optimal transport. Journal of Mathematical Imaging and Vision , 56(2):238-259, 2016.

N. Segol and Y. Lipman. On universal equivariant set networks. arXiv preprint arXiv:1910.02421 , 2019.

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.

A. Taghvaei and A. Jalali. 2-wasserstein approximation via restricted convex potentials with application to improved training for gans. arXiv preprint arXiv:1902.07197 , 2019.

M. A. Uy, Q.-H. Pham, B.-S. Hua, T. Nguyen, and S.-K. Yeung. Revisiting point cloud classification: A new benchmark dataset and classification model on real-world data. In Proceedings of the IEEE/CVF international conference on computer vision , pages 1588-1597, 2019.

E. Wagstaff, F. Fuchs, M. Engelcke, I. Posner, and M. A. Osborne. On the limitations of representing functions on sets. In International Conference on Machine Learning , pages 6487-6494. PMLR, 2019.

E. Wagstaff, F. B. Fuchs, M. Engelcke, M. A. Osborne, and I. Posner. Universal approximation of functions on sets. Journal of Machine Learning Research , 23(151):1-56, 2022.

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-1920, 2015.

Z. Yao, C. T. van Velthoven, T. N. Nguyen, J. 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-3241, 2021.

D. Yarotsky. Universal approximations of invariant maps by neural networks. Constructive Approximation , 55(1):407-474, 2022.

M. Zaheer, S. Kottur, S. Ravanbakhsh, B. Poczos, R. R. Salakhutdinov, and A. J. Smola. Deep sets. Advances in neural information processing systems , 30, 2017.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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\nSTRENGTHS:\nThe 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\n\nWEAKNESSES:\nThe 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\nQUESTIONS:\nSuggestions:\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\n\nLIMITATIONS:\nExponential dependence on dimension of points in the set should be mentioned.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-06-21"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n1. *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\nWEAKNESSES:\n1. *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.\n\nQUESTIONS:\nSee Section 'Weaknesses'.\n\nLIMITATIONS:\nThe authors should devote more space to the limitations of the proposed approach.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-04"}

comments: 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.

This paper presents a general neural network architecture for approximating symmetric and factor-wise group invariant (SFGI) functions.

The proposed SFGI neural network could achieve universality with the number of parameters independent of the input point set size.

The experiment results show the proposed SFGI neural network could approximate the (best or comparable) Wasserstein distance between point sets.

STRENGTHS:
This paper presents a general neural network framework for approximating the Wasserstein distance between point sets. 

The 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.

The proposed neural network approximator is simple but may inspire dozens of downstream network designs.

The presentation is clear. Well written.

The code is provided.

I 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)?

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.

Any insights for different distances and different complex objects (e.g. graph)? Should the proposed model to change accordingly?

LIMITATIONS:
NA
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2023-07-07

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nThe 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.\n\nQUESTIONS:\n1. 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\nLIMITATIONS:\nThe paper has no immediate negative societal impact.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}