SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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  $(\delta, 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  $(\Omega, \mathrm{d}_{\Omega})$ , let  $\mathcal{X}$  be the set of weighted point sets with at most  $N$  elements. In other words, each  $S \in \mathcal{X}$  has the form  $S = \{(x_i, w_i) : w_i \in \mathbb{R}_+, \sum_i w_i = 1, x_i \in \Omega\}$  and  $|S| \leq N$ . One can also consider  $\mathcal{X}$  to be the set of weighted Dirac measures over  $\Omega$ . 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  $\mathcal{X}$  to be the  $p$ -th Wasserstein distance,  $\mathrm{W}_p$ . We refer to the metric space of weighted point sets over  $\Omega$  as  $(\mathcal{X}, \mathrm{W}_p)$ .
First, we will show that given a  $\delta$ , there is a  $(\delta, a, G)$ -sketch of  $\mathcal{X}$  with respect to  $W_p$ . The embedding dimension  $a$  depends on the so-called covering number of the metric space  $(\Omega, \mathrm{d}_{\Omega})$  from which points are sampled. Given a compact metric space  $(\Omega, \mathrm{d}_{\Omega})$ , the covering number  $\nu_{\Omega}(r)$  w.r.t. radius  $r$  is the minimal number of radius  $r$  balls needed to cover  $\Omega$ . As a simple example, consider  $\Omega = [-\Delta, \Delta] \subseteq \mathbb{R}$ . Given any  $r$ , we can cover  $X$  with  $\frac{2\Delta}{r}$  intervals so  $\nu_{\Omega}(r) \leq \frac{2\Delta}{r}$ . The collection of the center of a set of  $r$ -balls that cover  $\Omega$  an  $r$ -net of  $\Omega$ . For a compact set  $\Omega \subset \mathbb{R}^d$  with diameter  $D$ , its covering number  $\nu_{\Omega}(r)$  is a constant depending on  $D$ ,  $r$  and  $d$  only.
Theorem 3.4 Set  $\mathrm{d}\chi$  to be  $W_{p}$  for  $1\leq p < \infty$ . Let  $G$  be the permutation group. For any  $\delta >0$ , let  $\delta_0 = \frac{1}{2}\sqrt[p]{\delta / 2}$  and  $a = \nu_{\Omega}(\delta_0)$  be the covering number w.r.t. radius  $\delta_0$ . Then there is a  $(\delta ,a,G)$ -sketch of  $\mathcal{X}$  with respect to  $\mathrm{W}_p$ . Furthermore, the encoding function  $\mathbf{h}:\mathcal{X}\to \mathbb{R}^a$  can be expressed as the following where  $h:\Omega \rightarrow \mathbb{R}^a$  is continuous:
$$
\mathbf {h} (S) = \sum_ {x \in S} h (x). \tag {1}
$$
Proof: Let  $\delta > 0$  and let  $S \in \mathcal{X}$  be  $S = \{(x_i, w_i) : \sum w_i = 1, x_i \in \Omega\}$  and  $|S| \leq N$ . Given  $\delta_0 = \frac{1}{2} \sqrt[p]{\delta/2}$  and  $a = \nu_{\Omega}(\delta_0)$ , we know  $\Omega$  has a  $\delta_0$ -net,  $C$ , and we denote the elements of  $C$  as  $\{y_1, \ldots, y_a\}$ . In other words, for any  $x \in \Omega$ , there is a  $y_i \in Cd$  such that  $\mathrm{d}_{\Omega}(x, y_i) < \delta_0$ .
First, we will define an encoding function  $\rho : \mathcal{X} \to \mathbb{R}^a$ . For each  $y_i$ , we will use a soft indicator function  $e^{-b\mathrm{d}_{\Omega}(x,B_{\delta_0}(y_i))}$  and set the constant  $b$  so that  $e^{-b\mathrm{d}_{\Omega}(x,B_{\delta_0}(y_i))}$  is "sufficiently" small if  $\mathrm{d}_{\Omega}(x,B_{\delta_0}(y_i)) > \delta_0$ . More formally, we know that  $\lim_{b \to \infty} e^{-b\delta_0} = 0$  so there is  $\beta \in \mathbb{R}$  such that for all  $b > \beta$ ,  $e^{-b\delta_0} < \frac{\delta_0^p}{d_{max}^p \cdot a}$ . Set  $b_0$  to be such that  $b_0 > \beta$ . Let  $h_i(x) = e^{-b_0\mathrm{d}_{\Omega}(x,B_{\delta_0}(y_i))}$  for each  $i \in [a]$ . For a given  $x \in \Omega$ , we compute  $h : \Omega \to \mathbb{R}^a$  as
$$
h (x) = \left[ h _ {1} (x), \dots , h _ {a} (x) \right]
$$
Then we define the encoding function  $\mathbf{h}:\mathcal{X}\to \mathbb{R}^a$  as
$$
\mathbf {h} (S) = \sum_ {i = 1} ^ {n} w _ {i} \frac {h \left(x _ {i}\right)}{\| h \left(x _ {i}\right) \| _ {1}}
$$
Note that  $\| \mathbf{h}(S)\| _1 = 1$  and  $\mathbf{h}$  is continuous since Wasserstein distances metrize weak convergence. Additionally, since  $\mathrm{d}_{\Omega}(x,B_{\delta_0}(y_i))$  is the distance from  $x$  to the  $\delta_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: \mathbb{R}^a \to \mathcal{X}$  as  $g(v) = \{(y_i, \frac{v_i}{\|v\|_1}) : i \in [a]\}$ . In order to show that  $g$  and  $\mathbf{h}$  yields a valid  $(\delta, a, G)$ -sketch of  $\mathcal{X}$ , we must show that  $g \circ \mathbf{h}(S)$  is sufficiently close to  $S$  under the  $\mathrm{W}_p$  distance. First, we know that
$$
\mathrm {W} _ {p} ^ {p} (g \circ \mathbf {h} (S), S) \leq \sum_ {i = 1} ^ {n} \sum_ {j = 1} ^ {a} w _ {1} \frac {h _ {j} (x _ {i})}{\| h (x _ {i}) \| _ {1}} d (x _ {i}, y _ {j}) ^ {p}.
$$
Let  $d_{max}$  be the diameter of  $\Omega$ . For a given  $x_i$ , we can partition  $\{h_1(x_i), \ldots, h_a(x_i)\}$  into those where  $h_j(x_i) \geq \frac{\delta_0^p}{d_{max}^p \cdot a}$  and those where  $h_j(x_i) < \frac{\delta_0^p}{d_{max}^p \cdot a}$  i.e.  $\{h_{j_1}(x_i), \ldots, h_{j_k}(x_i)\}$  and  $\{h_{j_{k+1}}(x_i), \ldots, h_{j_a}(x_i)\}$  respectively. If  $h_j(x) \geq \frac{\delta_0^p}{d_{max}^p \cdot a}$ , then
$$
e ^ {- b _ {0} \mathrm {d} _ {\Omega} (x, B _ {\delta_ {0}} (y _ {i}))} \geq \frac {\delta_ {0} ^ {p}}{d _ {m a x} ^ {p} \cdot a} > e ^ {- b _ {0} \delta_ {0}}
$$
so  $\mathrm{d}_{\Omega}(x,B_{\delta_0}(y_i)) <   \delta_0$  . Then
$$
\begin{array}{l} \mathrm {W} _ {p} ^ {p} (g \circ \mathbf {h} (S), S) \leq \sum_ {i = 1} ^ {n} \sum_ {j = 1} ^ {m} w _ {i} \frac {h _ {j} \left(x _ {i}\right)}{\| h \left(x _ {i}\right) \| _ {1}} \mathrm {d} _ {\Omega} \left(x _ {i}, y _ {j}\right) ^ {p}. \\ <   \sum_ {i = 1} ^ {n} w _ {i} \left(\sum_ {\ell = 1} ^ {k} \frac {h _ {j _ {\ell}} (x _ {i})}{\| h (x _ {i}) \| _ {1}} (2 \delta_ {0}) ^ {p} + \sum_ {\ell = k + 1} ^ {a} \frac {\delta_ {0} ^ {p}}{d _ {m a x} ^ {p} \cdot a} d _ {m a x} ^ {p}\right) \text {s i n c e} \mathrm {d} _ {\Omega} (x _ {i}, y _ {j}) \leq d _ {m a x} \\ \leq \sum_ {i = 1} ^ {n} w _ {i} \left(2 ^ {p} \delta_ {0} ^ {p} + \delta_ {0} ^ {p}\right) \leq 2 ^ {p} \left(\sqrt [ p ]{\delta / 2} \cdot \frac {1}{2}\right) ^ {p} + \frac {1}{2 ^ {p}} \cdot \frac {\delta}{2} <   \frac {\delta}{2} + \frac {\delta}{2} = \delta \\ \end{array}
$$
Thus, the encoding function  $\mathbf{h}$  and the decoding function  $g$  make up a  $(\delta, 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  $\epsilon$ -approximation of  $f$  via sum-pooling of continuous functions.
Corollary 3.5 Let  $\epsilon >0$ ,  $(\Omega ,d_{\Omega})$  be a compact metric space and let  $\mathcal{X}$  be the space of weighted point sets equipped with the  $p$ -Wasserstein,  $\mathrm{W}_p$ . Suppose for any  $\delta$ ,  $(\Omega ,d_{\Omega})$  has covering number  $a(\delta)$ . Then given a function  $f:\mathcal{X}\times \mathcal{X}\to \mathbb{R}$  that is uniformly continuous and permutation invariant, there is continuous functions  $h:\Omega \to \mathbb{R}^{a(\delta)}$ ,  $\phi :\mathbb{R}^{a(\delta)}\to \mathbb{R}^{a'}$ , and  $\rho :\mathbb{R}^{a^{\prime}}\to \mathbb{R}$ , such that for any  $A,B\in \mathcal{X}$
$$
\left| f (A, B) - \rho \left(\phi \left(\sum_ {(x, w _ {x}) \in A} w _ {x} h (x)\right) + \phi \left(\sum_ {(x, w _ {x}) \in B} w _ {x} h (x)\right)\right) \right| <   \epsilon
$$
where  $h, \phi$  and  $\rho$  are all continuous and  $a' = 4 \cdot a(\delta) + 1$ .
Due to Eqn. (1), instead of considering the function  $\mathbf{h}$  which takes a set of points  $S \in \mathcal{X}$  as input, we now only need to model the function  $h: \Omega \to \mathbb{R}^a$ , which takes a single point  $x \in S$  as input. For simplicity, assume that the input metric space  $(\Omega, d_{\Omega})$  is a compact set in some Euclidean space  $\mathbb{R}^d$ . Note that in contrast to Lemma 3.3, each  $h, \phi$  and  $\rho$  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  $\Omega$ , the target additive error  $\epsilon$ , and independent of the maximum size  $N$  of input points. We can use multilayer perceptrons (MLPs)  $h_{\theta_1}, \phi_{\theta_2}$ , and  $\rho_{\theta_3}$  in place of  $h, \phi$  and  $\rho$  to approximate the desired function. Formally, we define the following family of neural networks:
$$
\mathcal {N} _ {\text {P r o d u c t N e t}} (A, B) = \rho_ {\theta_ {3}} \left(\phi_ {\theta_ {2}} \left(\sum_ {(x, w _ {x}) \in A} w _ {x} h _ {\theta_ {1}} (x)\right) + \phi_ {\theta_ {2}} \left(\sum_ {(x, w _ {x}) \in B} w _ {x} h _ {\theta_ {1}} (x)\right)\right). \tag {2}
$$
In practice, we consider the input to the neural network  $h_{\theta_1}$  to be a point  $x \in \Omega$  along with its weight  $w_x$ . As per the discussions above, functions represented by  $\mathcal{N}_{\mathrm{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,  $\mathrm{W}_p: \mathcal{X} \times \mathcal{X} \to \mathbb{R}$ , are uniformly continuous with respect to the underlying metric  $\mathrm{W}_p$ , we can apply our framework for the problem of approximating  $p$ -Wasserstein.
Importantly, the number of parameters in  $\mathcal{N}_{\mathrm{ProductNet}}$  does not depend on the maximum size of the point set but rather only on the  $\epsilon$  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  $\mathbb{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  $\rho_{\theta_3}$ ,  $\phi_{\theta_2}$ , and  $h_{\theta_1}$  depend on the dimensions of the domain and co-domain of each function they approximate ( $\rho$ ,  $\phi$  and  $h$ ) and the desired approximation error  $\epsilon$ . We assume that MLPs are Lipschitz continuous. In our case,  $\phi_{\theta_2}$  operates on the sum of  $h_{\theta_1}(x)$ s for all  $N$  numbers of input points in  $A$  (or in  $B$ ). In general, the error made in  $hs$  may accumulate  $N$  times, which causes the precision we must achieve for each individual  $h_{\theta_1}(x)$  (compared to  $h(x)$ ) to be  $\Theta(\epsilon/N)$ . This would have caused the model complexity of  $h_{\theta_1}$  to depend on  $N$ . Fortunately, this is not the case for our encoding
Figure 1: Visually representing a neural network which can universally approximate uniformly continuous SFGI product functions over pairs of point sets.
function  $h$ . In particular, because our encoding function can be written in the specific form of a normalized sum of individual points in the set  $S$ ; i.e.,  $\phi_{\theta_2}$  operates on  $\sum_{(x, w_x) \in S} w_x h_{\theta_1}(x)$  with  $\sum_{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 \in S$ . Thus, as both the error for each MLP and the dimension of the domain and co-domain of each approximated function  $(\rho, \phi, h)$  do not depend on the size of the input point set  $N$ , we get that  $\rho_{\theta_3}, \phi_{\theta_2}$  and  $h_{\theta_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  $\mathcal{N}_{\mathrm{ProductNet}}$  which can approximate the  $p$ -Wasserstein distance between point sets to within an additive  $\epsilon$ -error. The number of parameters in this network depends on  $\epsilon$  and the covering number of  $\Omega$  and not on the size of each point set.
Additionally, if we replace the sum-pooling with max in  $\mathcal{N}_{\mathrm{ProductNet}}$ , there is a network of such a form where we can also approximate Hausdorff distance between point sets to additive  $\epsilon$  accuracy.
Exponential dependence on dimension Although the model complexity of  $\mathcal{N}_{\mathrm{ProductNet}}$  is independent of the size of the input point set, it depends on the covering number of  $\Omega$  which, in turn, can have an exponential dependence on the dimension of  $\Omega$ . In short, this means that the model complexity of  $\mathcal{N}_{\mathrm{ProductNet}}$  has an exponential dependence on the dimension of  $\Omega$ . 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  $\Omega$  will depend on the intrinsic dimension of  $\Omega$  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 \in \mathcal{X}$ ,  $\sum_{(x, w_x) \in A} w_x h_{\theta_1}(x) + \sum_{(x, w_x) \in B} w_x h_{\theta_1}(x)$ , we use the sum of these two components (as in a DeepSet architecture):  $\sum_{(x, w_x) \in A} w_x h_{\theta_1}(x)$  and  $\sum_{(x, w_x) \in B} w_x h_{\theta_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 \to \mathbb{R}$  where  $K \subseteq \mathbb{R}^{a \times 2}$  is compact, we need
Table 1: Mean relative error between approximations and ground truth Wasserstein distance between point sets. The top row for each dataset shows the approximation quality for point sets with input sizes that were seen at training time; while he bottom row shows the approximation quality for point sets with input sizes that were not seen at training time. Note that  $\mathcal{N}_{\mathrm{ProductNet}}$  is our model.
Dataset	Input size	NProductNet	WPCE	NSDeepSets	Sinkhorn
noisy-sphere-3	[100, 300]	0.046 ± 0.043	0.341 ± 0.202	0.362 ± 0.241	0.187 ± 0.232
[300, 500]	0.158 ± 0.198	0.356 ± 0.286	0.608 ± 0.431	0.241 ± 0.325
noisy-sphere-6	[100, 300]	0.015 ± 0.014	0.269 ± 0.285	0.291 ± 0.316	0.137 ± 0.122
[300, 500]	0.049 ± 0.054	0.423 ± 0.408	0.508 ± 0.473	0.198 ± 0.181
uniform	256	0.097 ± 0.073	0.120 ± 0.103	0.123 ± 0.092	0.073 ± 0.009
[200, 300]	0.131 ± 0.096	1.712 ± 0.869	0.917 ± 0.869	0.064 ± 0.008
ModelNet-small	[20, 200]	0.084 ± 0.077	0.077 ± 0.075	0.105 ± 0.096	0.101 ± 0.032
[300, 500]	0.111 ± 0.086	0.241 ± 0.198	0.261 ± 0.245	0.193 ± 0.155
ModelNet-large	2048	0.140 ± 0.206	0.159 ± 0.141	0.166 ± 0.129	0.148 ± 0.048
[1800, 2000]	0.162 ± 0.228	0.392 ± 0.378	0.470 ± 0.628	0.188 ± 0.088
RNAseq	[20, 200]	0.012 ± 0.010	0.477 ± 0.281	0.482 ± 0.291	0.040 ± 0.009
[300, 500]	0.292 ± 0.041	0.583 ± 0.309	0.575 ± 0.302	0.048 ± 0.006
Table 2: Training time and number of epochs needed for convergence for best model
Dataset	NProductNet	WPCE	NSDeepSets
noisy-sphere-3	Time	6min	1hr 46min	9min
Epochs	20	100	100
noisy-sphere-6	Time	12min	4hr 6min	1hr 38min
Epochs	20	100	100
uniform	Time	7min	3hr 36min	1hr 27min
Epochs	23	100	100
ModelNet-small	Time	7min	1hr 23min	12 min
Epochs	20	100	100
ModelNet-large	Time	8min	3hr 5min	40min
Epochs	20	100	100
RNAseq(2k)	Time	15min	14hr 26min	3hr 01min
Epochs	73	100	100
to use a  $\delta$ -net for  $K$  which will also serve as the intermediate dimension for  $\phi$ . As  $K \subseteq [0, N]^{a \times 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.
We evaluate the accuracy of the 1-Wasserstein distance approximations of our proposed neural network architecture,  $\mathcal{N}_{\mathrm{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  $\mathcal{N}_{\mathrm{SDep}\mathrm{eps}}$ , which is a single DeepSets model which maps both point sets to a Euclidean space and approximates the 1-Wasserstein distance as the  $\ell_2$  norm between output of each point set. As Siamese networks are widely employed for metric learning,  $\mathcal{N}_{\mathrm{SDep}\mathrm{eps}}$  model is a natural baseline for comparison again  $\mathcal{N}_{\mathrm{ProductNet}}$ . We additionally test our neural network approximations against the Sinkhorn distance where the regularization parameter was set to  $\epsilon = 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  $\mathbb{R}^3$  and a gene expression dataset (RNAseq) which consists of 4360 cells each represented by 2000 genes (i.e. 4360 points in  $\mathbb{R}^{2000}$ ) as well as three synthetic datasets: (1) uniform, where point sets are in  $\mathbb{R}^2$ , (2) noisy-sphere-3, where point sets are in  $\mathbb{R}^3$ , (3) noisy-sphere-6, where point sets are in  $\mathbb{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,  $\mathcal{N}_{\mathrm{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  $\mathcal{N}_{\mathrm{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),  $\mathcal{N}_{\mathrm{ProductNet}}$  showed significantly lower error than all other methods including WPCE. These results along with a more detailed plot in Figure 2 in Appendix D indicate that  $\mathcal{N}_{\mathrm{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  $\mathcal{N}_{\mathrm{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,  $\mathcal{N}_{\mathrm{ProductNet}}$  can reduce training time while still achieving comparable or better quality approximations of Wasserstein distance. Using our  $\mathcal{N}_{\mathrm{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.
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  $\mathcal{N}_{\mathrm{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 desirable 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  $\epsilon$ -error where the complexity does not depend on the maximum size of the input set.
The authors thank anonymous reviewers for their helpful comments. Furthermore, the authors thank Rohan Gala and the Allen Institute for generously providing the RNAseq dataset. Finally, Samantha Chen would like to thank Puoya Tabaghi for many helpful discussions about permutation invariant neural networks and Tristan Brugère for his implementation of the Sinkhorn algorithm. This research is supported by National Science Foundation (NSF) grants CCF-2112665 and CCF-2217058 and National Institutes of Health (NIH) grant RF1 MH125317.