Permutation invariant networks to learn Wasserstein metrics

Abstract

Understanding the space of probability measures on a metric space equipped with a Wasserstein distance is one of the fundamental questions in mathematical analysis. The Wasserstein metric has received a lot of attention in the machine learning community especially for its principled way of comparing distributions. In this work, we use a permutation invariant network to map samples from probability measures into a low-dimensional space such that the Euclidean distance between the encoded samples reflects the Wasserstein distance between probability measures. We show that our network can generalize to correctly compute distances between unseen densities. We also show that these networks can learn the first and the second moments of probability distributions.

arXiv commentFix typos, Accepted as a spotlight at Topological Data Analysis and Beyond Workshop at Neurips 2020. Added more experiments and results. Comments welcome

References in corpus (9)

• Fast and Accurate Deep Network Learning by Exponential Linear Units (ELUs)
• Statistical Aspects of Wasserstein Distances
• Fast Computation of Wasserstein Barycenters
• Learning with a Wasserstein Loss
• Near-linear time approximation algorithms for optimal transport via Sinkhorn iteration
• Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distances
• Riemannian Center of Mass and so called karcher mean
• Learning Probability Measures with respect to Optimal Transport Metrics
• Discrete Wasserstein Barycenters: Optimal Transport for Discrete Data