Neural approximation of Wasserstein distance via a universal architecture for symmetric and factorwise group invariant functions

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.

Strengths

The paper presents two nice ideas: (a) Characterizaton of SFGI functions (not so suprising in my humble opinion, but still new) (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. Additionally 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.

Questions

Suggestions: Explain a bit more re why to learn Wasserstien distances. Mention you results in the appendix which show better inference time than Sinkhorn. In 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. Line 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). Definition of W_p: D should be replaced by D^p. Also the index of w_n' should be w_m'. Definition 3.1: specify the group G acts on X. Lemma 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? Line 185: `then we have that'--> then we can choose Line 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? Line 208: i.e., MLP should be e.g. MLP Line 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 I would view the construction in Theorem 3.4 as a smoothened version of one-hot encoding. You might want to mention this connection. Line 238: M was N previously, better to be consistent. Line 244: I'm not sure whether d_max was defined? (presumably it is the diameter of X) The 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)). The second line in said inequality has an explanation in the middle, this could be styled better. In 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. Corollary 3.6: mention dependence on dimension not only on epsilon. Line 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 [1] Low Dimensional Invariant Embeddings for Universal Geometric Learning, https://arxiv.org/pdf/2205.02956.pdf [2] Neural Injective Functions for Multisets, Measures and Graphs via a Finite Witness Theorem https://arxiv.org/pdf/2306.06529.pdf

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

5: You are absolutely certain about your assessment. You are very familiar with the related work and checked the math/other details carefully.

Soundness

4 excellent

Presentation

3 good

Contribution

3 good

Limitations

Exponential dependence on dimension of points in the set should be mentioned.

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.

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. 2. *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. 3. *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. 4. *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.

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. 2. *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.

Questions

See Section 'Weaknesses'.

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

The authors should devote more space to the limitations of the proposed approach.

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?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

4 excellent

Presentation

4 excellent

Contribution

4 excellent

Limitations

NA

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.

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.

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. 2. Shouldn't it be $D_{i,j}=d_{X}(x_{i},x'_{j})^{p}$ instead in Line 153? 3. 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). 4. 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$?

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

The paper has no immediate negative societal impact.

Response to the rebuttal

I appreciate the detailed response by the authors. Taking into account the other reviews and corresponding rebuttals, I stick to my initial evaluation.