SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Samantha Chen  
Department of Computer Science Engineering  
University of California - San Diego  
sac003@ucsd.edu
Yusu Wang  
Halicioglu Data Science Institute  
University of California - San Diego  
yusuwang@ucsd.edu
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).
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]
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  $\mathrm{W}_p(X,Y)$  between two point sets  $X$  and  $Y$  sampled from some metric space (e.g.,  $\mathbb{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,\ldots ,P_k$  in  $\mathbb{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  $\mathcal{X}$  denote the space of finite point sets from a bounded region in  $\mathbb{R}^d$ . Then the  $p$ -Wasserstein distance can be viewed as a function  $\mathrm{W}_p: \mathcal{X} \times \mathcal{X} \to \mathbb{R}$ . Similarly, 1-median for  $k$  point sets can be modeled by a function from the product of  $k$  copies of  $\mathcal{X}$  to  $\mathbb{R}$ . Such functions are not only invariant to permutations of the factors of the product space (i.e.  $\mathrm{W}_p(X, Y) = \mathrm{W}_p(Y, X)$ ) but are also invariant or equivariant with respect to certain group actions for each factor. For  $p$ -Wasserstein distance,  $\mathrm{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  $(\mathcal{X}_i,\mathrm{d}_{\mathcal{X}_i})$  where  $i\in [k]$ , we define symmetric and factor-wise group invariant (SFGI) product functions as  $f:\mathcal{X}_1\times \mathcal{X}_2\times \dots \mathcal{X}_k\to \mathbb{R}$  where  $f$  is (1) symmetric over the  $k$  factors, and (2) invariant to the group action  $G_{1}\times G_{2}\times \dots \times G_{k}$  for some group  $G_{i}$  acting on  $\mathcal{X}_i$ , for each  $i\in [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 thresholds. 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 2-Wasserstein 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  $\epsilon$ -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  $\ell_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  $\ell_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.
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  $(\mathcal{X},\mathrm{d}_{\mathcal{X}})$ , the action of the group  $G$  on  $\mathcal{X}$  is a function  $\alpha : G\times \mathcal{X}\to \mathcal{X}$  that transforms the elements of  $\mathcal{X}$  for each element  $\pi \in G$ . For each element  $\pi \in G$ , we will write  $\pi \cdot x$  to denote the action of a group element  $\pi$  on  $x$  instead of  $\alpha (\pi ,x)$ . For example, if  $G$  is the permutation group over  $[N]\coloneqq \{1,2,\ldots ,N\}$ , and  $\mathcal{X} = \mathbb{R}^N$ , then for any  $\pi \in G$ ,  $\pi \cdot x$  represents the permutation of elements in  $x\in \mathbb{R}^N$  via  $\pi$  i.e. given  $x = (x_{1},x_{2},\dots,x_{N})$ ,  $\pi \cdot x = (x_{\pi (1)},x_{\pi (2)}\dots,x_{\pi (N)})$ . A function  $f:\mathcal{X}\rightarrow \mathcal{Y}$  is  $G$ -invariant if for any  $X\in \mathcal{X}$  and any  $\pi \in G$ , we have that  $f(X) = f(\pi \cdot X)$ .
Universal Approximation. Let  $\mathcal{C}(\mathcal{X},\mathbb{R})$  denote the set of continuous functions from a metric space  $(\mathcal{X},d_{\mathcal{X}})$  to  $\mathbb{R}$ . Given two families of functions  $\mathcal{F}_1$  and  $\mathcal{F}_2$  where  $\mathcal{F}_1\subseteq \mathcal{F}_2$  and  $\mathcal{F}_1,\mathcal{F}_2\subseteq \mathcal{C}(\mathcal{X},\mathbb{R})$ , we say that  $\mathcal{F}_1$  universally approximates  $\mathcal{F}_2$  if for any  $\epsilon >0$  and any  $f\in \mathcal{F}_2$ , there is a  $g\in \mathcal{F}_1$  such that  $\| g - f\|_{\infty} < \epsilon$ . Different norms on the space of functions can be used, but we will use  $L_{\infty}$  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 output embedding. Formally, suppose we are given a finite multiset  $S = \{x : x \in \mathbb{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
$$
\mathcal {N} _ {\text {D e e p S e t}} (S) = g _ {\theta_ {2}} \left(\sum_ {x \in S} h _ {\theta_ {1}} (x)\right)
$$
where  $h_{\theta_1}$  and  $g_{\theta_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  $\mathbb{R}^k$ , and the input multiset size is fixed as  $N$ . Let  $t = 2kN + 1$ . Then any continuous multiset function, represented as  $f: \mathbb{R}^{k \times N} \to \mathbb{R}$  which is invariant with respect to permutations of the columns, can be approximated arbitrarily close in the form of  $\rho\left(\sum_{x \in X} \phi(x)\right)$ , for continuous transformations  $\phi: \mathbb{R}^k \to \mathbb{R}^t$  and  $\rho: \mathbb{R}^t \to \mathbb{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  $\phi$  is  $t = \binom{k+N}{k}$ . It was recently shown in [12, 1] that the output dimension of  $\phi$  can be reduced to  $2kN + 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  $\phi$  is  $t = \binom{k+N}{k}$  or  $t = 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, \mathrm{d}_X)$  be a metric space. For two weighted point sets  $P = \{(x_i, w_i) : x_i \in X, \sum_{w_i} = 1, i \in [n]\}$  and  $Q = \{(x_i', w_i') : x_i' \in X, \sum_{w_i} = 1, i \in [m]\}$ , we define the Wasserstein distance between  $P$  and  $Q$  as
$$
\mathrm {W} _ {p} (P, Q) = \min _ {\Pi \in \mathbb {R} _ {+} ^ {n \times m}} \left\{\left(\langle \Pi , D ^ {p} \rangle\right) ^ {1 / p}: \Pi \mathbf {1} = [ w _ {1}, \ldots , w _ {n} ], \Pi^ {T} \mathbf {1} = [ w _ {1} ^ {\prime}, \ldots , w _ {m} ^ {\prime} ] \right\}
$$
where  $D \in \mathbb{R}_+^{n \times m}$  is the distance matrix with  $D_{i,j} = \mathrm{d}_X(x_i, x_j')$ . One can think of these weighted point sets as discrete probability distributions in  $(X, \mathrm{d}_X)$ . When  $p = 1$ ,  $\mathrm{W}_1$  is also commonly known as the Earth Mover's distance (EMD). Additionally, note that when  $p = \infty$ ,  $\mathrm{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  $\epsilon$ . As  $\epsilon$  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  $\delta >0$ ,  $a\in \mathbb{N}^+$ , and  $G$  be a group which acts on  $\mathcal{X}$ . A  $(\delta ,a,G)$ -sketch of a metric space  $(\mathcal{X},\mathrm{d}_{\mathcal{X}})$  consists of a  $G$ -invariant continuous encoding function  $h:\mathcal{X}\to \mathbb{R}^a$  and a continuous decoding function  $g:\mathbb{R}^a\rightarrow \mathcal{X}$  such that  $\mathrm{d}_{\mathcal{X}}(g\circ h(S),S) < \delta$ .
Now let  $(\mathcal{X}_1, \mathrm{d}_{\mathcal{X}_1}), \ldots, (\mathcal{X}_m, \mathrm{d}_{\mathcal{X}_m})$  be compact metric spaces. The product space  $\mathcal{X}_1 \times \dots \times \mathcal{X}_m$  is still a metric space equipped with the following natural metric induced from metrics of each factor:
$$
\mathrm {d} _ {\mathcal {X} _ {1} \times \dots \times \mathcal {X} _ {m}} ((A _ {1}, \dots , A _ {m}), (A _ {1} ^ {\prime}, \dots , A _ {m} ^ {\prime})) = \mathrm {d} _ {\mathcal {X} _ {1}} (A _ {1}, A _ {1} ^ {\prime}) + \dots + \mathrm {d} _ {\mathcal {X} _ {m}} (A _ {m}, A _ {m} ^ {\prime}).
$$
Suppose  $G_{i}$  is a group acting on  $\mathcal{X}_i$ , for each  $i \in [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: \mathcal{X}_1 \times \dots \times \mathcal{X}_m \to \mathbb{R}$  such that  $f$  is uniformly continuous and invariant to the group action  $G_{1} \times \dots \times G_{m}$ .
Lemma 3.2 Suppose  $f: \mathcal{X}_1 \times \dots \times \mathcal{X}_m \to \mathbb{R}$  is uniformly continuous and invariant to  $G_1 \times \dots \times G_m$ . Additionally, assume that for any  $\delta > 0$ ,  $(\mathcal{X}_i, \mathrm{d}_{\mathcal{X}_i})$  has a  $(\delta, a_i, G_i)$ -sketch where  $a_i$  may depend on  $\delta$ . Then for any  $\epsilon > 0$ , there is a continuous  $G_i$ -invariant function  $\phi_i: \mathcal{X}_i \to \mathbb{R}^{a_i}$  for all  $i \in [k]$  and a continuous function  $\rho: \mathbb{R}^{a_1} \times \dots \times \mathbb{R}^{a_m} \to \mathbb{R}$  such that
$$
\left| f \left(A _ {1}, A _ {2}, \dots , A _ {m}\right) - \rho \left(\phi_ {1} \left(A _ {1}\right), \phi_ {2} \left(A _ {2}\right), \dots , \phi_ {k} \left(A _ {m}\right)\right) \right| <   \epsilon
$$
for any  $(A_{1},\ldots ,A_{m})\in \mathcal{X}_{1}\times \dots \times \mathcal{X}_{m}$ . Furthermore, if  $\mathcal{X}_1 = \dots \mathcal{X}_2 = \dots = \mathcal{X}_m$ , then we can choose  $\phi_1 = \phi_2 = \dots \phi_m$
Note that a recent result from [17] shows that a continuous factor-wise group invariant function  $f: \mathcal{X}_1 \times \dots \times \mathcal{X}_m \to \mathbb{R}$  can be represented (not approximated) by the form  $f(v_1, \ldots, v_m) = \rho(\phi_1(v_1), \ldots, \phi_k(v_m))$  if there exists a topological embedding from  $\mathcal{X}_i / G_i$  to Euclidean space. The condition that each quotient  $\mathcal{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  $(\delta, a, G)$ -sketch).
We often consider the case where  $\mathcal{X} = \mathcal{X}_1 = \dots = \mathcal{X}_m$  i.e.  $f:\mathcal{X}\times \dots \mathcal{X}\to \mathbb{R}$  where  $G$  is a group acting on the factor  $\mathcal{X}$ . Oftentimes, we require the function to not only be invariant to the actions of a group  $G$  on each individual  $\mathcal{X}$  but also symmetric with respect to the ordering of the input. By this, we mean  $f(A_{1},\ldots ,A_{m}) = f(A_{\pi (1)},\ldots ,A_{\pi (m)})$  where  $\pi$  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  $\mathcal{X} = \mathcal{X}_1 = \dots = \mathcal{X}_m$  and  $G = G_{1} = \dots = G_{m}$ . Assume that  $\mathcal{X}$  has a  $(\delta, a, G)$ -sketch. Additionally, suppose  $f$  is symmetric; hence  $f$  is a SFGI function. Let  $t = 2am + 1$ . Then for any  $\epsilon > 0$ , there is a continuous  $G$ -invariant function  $\phi: \mathcal{X} \to \mathbb{R}^t$  and a continuous function  $\rho: \mathbb{R}^t \to \mathbb{R}$  such that
$$
\left| f \left(A _ {1}, \dots , A _ {m}\right) - \rho \left(\sum_ {i = 1} ^ {m} \phi \left(A _ {i}\right)\right) \right| <   \epsilon
$$
Now suppose we want approximate an SFGI product function,  $f$ , with a neural network. Lemma 3.3 implies that we can approximate  $\phi$  with any universal  $G$ -invariant neural network which embeds our original space  $\mathcal{X}$  to some Euclidean space  $\mathbb{R}^a$ . Then the outer architecture  $\rho$  can be any universal architecture (e.g. MLP). Finding a universal  $G$ -invariant neural network to realize  $\phi$  over a single factor space  $\mathcal{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