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

arXiv: https://arxiv.org/abs/2308.00273 · HTML: https://mlsanity.com/paper/2308.00273/

## Key facts

- Cited by 4 papers in the map (data through 14 Oct 2026).
- Builds on 26 of its 42 references that are in the map.
- First posted on arXiv 1 Aug 2023.

## Builds on (most cited first, 26 of 26)

- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [PointNet: Deep Learning on Point Sets for 3D Classification and Segmentation](https://mlsanity.com/paper/1612.00593/) (2016)
- [Neural Message Passing for Quantum Chemistry](https://mlsanity.com/paper/1704.01212/) (2017)
- [Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distances](https://mlsanity.com/paper/1306.0895/) (2013)
- [Deep Sets](https://mlsanity.com/paper/1703.06114/) (2017)
- [Revisiting Point Cloud Classification: A New Benchmark Dataset and Classification Model on Real-World Data](https://mlsanity.com/paper/1908.04616/) (2019)
- [Invariant and Equivariant Graph Networks](https://mlsanity.com/paper/1812.09902/) (2018)
- [Vector Neurons: A General Framework for SO(3)-Equivariant Networks](https://mlsanity.com/paper/2104.12229/) (2021)
- [Optimal transport mapping via input convex neural networks](https://mlsanity.com/paper/1908.10962/) (2019)
- [Universal Invariant and Equivariant Graph Neural Networks](https://mlsanity.com/paper/1905.04943/) (2019)
- [Sign and Basis Invariant Networks for Spectral Graph Representation Learning](https://mlsanity.com/paper/2202.13013/) (2022)
- [On the Universality of Invariant Networks](https://mlsanity.com/paper/1901.09342/) (2019)
- [Universal approximations of invariant maps by neural networks](https://mlsanity.com/paper/1804.10306/) (2018)
- [On Learning Sets of Symmetric Elements](https://mlsanity.com/paper/2002.08599/) (2020)
- [On the Limitations of Representing Functions on Sets](https://mlsanity.com/paper/1901.09006/) (2019)
- [Molecular graph generation with Graph Neural Networks](https://mlsanity.com/paper/2012.07397/) (2020)
- [SE(3) Equivariant Graph Neural Networks with Complete Local Frames](https://mlsanity.com/paper/2110.14811/) (2021)
- [Universal Approximation of Functions on Sets](https://mlsanity.com/paper/2107.01959/) (2021)
- [On Universal Equivariant Set Networks](https://mlsanity.com/paper/1910.02421/) (2019)
- [Low Dimensional Invariant Embeddings for Universal Geometric Learning](https://mlsanity.com/paper/2205.02956/) (2022)
- [Neural Injective Functions for Multisets, Measures and Graphs via a Finite Witness Theorem](https://mlsanity.com/paper/2306.06529/) (2023)
- [2-Wasserstein Approximation via Restricted Convex Potentials with Application to Improved Training for GANs](https://mlsanity.com/paper/1902.07197/) (2019)
- [Learning Wasserstein Embeddings](https://mlsanity.com/paper/1710.07457/) (2017)
- [Scalable Nearest Neighbor Search for Optimal Transport](https://mlsanity.com/paper/1910.04126/) (2019)
- [A reconfigurable neural network ASIC for detector front-end data compression at the HL-LHC](https://mlsanity.com/paper/2105.01683/) (2021)
- [A Sparse Multi-Scale Algorithm for Dense Optimal Transport](https://mlsanity.com/paper/1510.05466/) (2015)

## Built on by (most cited first, 4 of 4)

- [Fourier Sliced-Wasserstein Embedding for Multisets and Measures](https://mlsanity.com/paper/2504.02544/) (2025)
- [Relative Translation Invariant Wasserstein Distance](https://mlsanity.com/paper/2409.02416/) (2024)
- [Revisiting Multi-Permutation Equivariance through the Lens of Irreducible Representations](https://mlsanity.com/paper/2410.06665/) (2024)
- [Fourier Sliced-Wasserstein Embedding for Multisets and Measures](https://mlsanity.com/paper/2405.16519/) (2024)

## Abstract (arXiv)

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