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 (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).

翻译：学习复杂对象之间的距离函数（如用于比较点集的Wasserstein距离）是机器学习应用中的常见目标。然而，此类复杂对象（例如点集和图）上的函数通常需要对多种群作用（如置换或刚体变换）具有不变性。因此，定义在此类复杂对象上的连续对称乘积函数（如距离函数）也必须对这类群作用的乘积具有不变性。我们称此类函数为对称与因子群不变函数（简称SFGI函数）。本文首先提出一种用于逼近SFGI函数的通用神经网络架构。本文的主要贡献在于，将这一通用神经网络与草图化思想相结合，开发出一种能够逼近点集之间$p$次Wasserstein距离的高效专用神经网络。极为重要的是，所需的模型复杂度与输入点集的大小无关。在理论层面，据我们所知，这是首个证明存在具有有界模型复杂度且能逼近Wasserstein距离的神经网络的研究成果。本工作为几何问题的草图化思想与对称函数通用逼近方法提供了有趣的融合。在实验层面，我们展示了一系列结果，表明新提出的神经网络架构性能与其他模型（包括基于最先进孪生自编码器的方法）相当或更优。特别是，我们的神经网络在泛化性能上显著更优，且训练速度远快于最先进的孪生自编码器。最后，这一研究方向可能有助于探索解决广泛几何优化问题（如度量空间中的$k$-均值聚类）的有效神经网络设计。