Closed-form analytical solutions to Optimal Transport/Wasserstein distance

5.4k Views Asked by At

Kuang and Tabak (2017) mentions that:

"closed-form solutions of the multidimensional optimal transport problems are relatively rare, a number of numerical algorithms have been proposed."

I'm wondering if there are some resources (lecture notes, papers, etc.) that collect/contain known solutions to optimal transport and/or Wasserstein distance between two distributions in dimensions greater than 1. For example, let $ \mathcal{N_1}(\mu_1, \Sigma_1) $ and $ \mathcal{N_2}(\mu_2, \Sigma_2) $ denote two Gaussian distributions with different means and covariances matrices. Then the optimal transport map between them is:

$$ x \longrightarrow \mu_2 + A( x - \mu_1 ) $$ where $ A = \Sigma_1^{- 1/2} (\Sigma_1^{1/2} \Sigma_2 \Sigma_1^{1/2})^{1/2} \Sigma_1^{- 1/2}$. And so the Wasserstein 2 distance is

$$ W_2 ( \mathcal{N_1}(\mu_1, \Sigma_1), \mathcal{N_2}(\mu_2, \Sigma_2) ) = || \mu_1 - \mu_2 ||^2_2 + \mathrm{Tr}( \Sigma_1 + \Sigma_2 - 2( \Sigma_1^{1/2} \Sigma_2 \Sigma_1^{1/2} )^{1/2} ) $$ where $\mathrm{Tr}$ is the trace operator.

It will be nice to know more worked out examples of optimal transport, such as uniform distributions between different geometric objects, e.g. concentric and overlapping balls, between rectangles, etc.

2

There are 2 best solutions below

4
On

Although a bit old, this is indeed a good question. Here is my bit on the matter:

  1. Regarding Gaussian Mixture Models, see A Wasserstein-type distance in the space of Gaussian Mixture Models, Julie Delon and Agnes Desolneux (2020).

  2. Using the 2-Wasserstein metric, Mallasto and Feragen geometrize the space of Gaussian processes with $L_2$ mean and covariance functions over compact index spaces: Learning from uncertain curves: The 2-Wasserstein metric for Gaussian processes, Anton Mallasto, Aasa Feragen https://papers.nips.cc/paper/7149-learning-from-uncertain-curves-the-2-wasserstein-metric-for-gaussian-processes.pdf

  3. Wasserstein space of elliptical distributions are characterized by Muzellec and Cuturi. Authors show that for elliptical probability distributions, Wasserstein distance can be computed via a simple Riemannian descent procedure: Generalizing Point Embeddings using the Wasserstein Space of Elliptical Distributions, Boris Muzellec and Marco Cuturi https://arxiv.org/pdf/1805.07594.pdf (Not closed form)

  4. Tree metrics as ground metrics yield negative definite OT metrics that can be computed in a closed form. Sliced-Wasserstein distance is then a particular (special) case (the tree is a chain): Tree-Sliced Variants of Wasserstein Distances, Tam Le, Makoto Yamada, Kenji Fukumizu, Marco Cuturi https://arxiv.org/pdf/1902.00342.pdf

  5. Sinkhorn distances/divergences (Cuturi, 2013) are now treated as new forms of distances (e.g. not approximations to $\mathcal{W}_2^2$) (Genevay et al, 2019). Recently, this entropy regularized optimal transport distance is found to admit a closed form for Gaussian measures: Janati et al (2020). This fascinating finding also extends to the unbalanced case.

  6. Variations of Wasserstein distance such as Gromov-Wasserstein distance have been studied for Gaussian distributions. For those cases, Salmona et al. (2019) have found a simple linear method for finding the Gromov-Monge map.

  7. Bunne et al. (2023) has shown that the dynamic formulation of optimal transport, also known as the Schrödinger bridge (SB) problem, admits a closed form solution for Gaussian measures. In other words, they provide closed-form expressions for SBs between Gaussian measures. This is exciting due to the connections to the recently popularized diffusion models.

  8. Chizat (2023) has introduced the $(\lambda, \tau)$-barycenter as the unique probability measure that minimizes the sum of entropic optimal transport. He has shown that for isotropic Gaussians, there is a closed-form solution.

  9. Bonet et al. (2023) have proposed a closed form solution for calculating the Wasserstein distance between an arbitrary distribution and a uniform distribution, both situated on a circular domain.

  10. Beraha and Pegoraro (2023) derive an (almost) closed form expression for optimal transportation of measures supported on the unit-circle.

I would be happy to keep this list up to date and evolving.

2
On

Optimal transport (OT) problems admit closed-form analytical solutions in a very few notable cases, e.g. in 1D or between Gaussians. Below I cite articles providing analytical solutions for the 1-dimensional case only (does 1D mean univariate?)

Formula 3 in the following gives a closed-form analytical solution for Wasserstein distance in the case of 1-D probability distributions, but a source for the formula isn't given and I wonder how to convert it to a discretized linear programming model:

  1. Kolouri et al (2019) "Generalized Sliced Wasserstein Distances" https://arxiv.org/pdf/1902.00434.pdf

Formula 9 in the following also gives a closed-form solution:

  1. Kolouri et al (2019) "Sliced-Wasserstein Auto-encoders" https://openreview.net/pdf?id=H1xaJn05FQ

Formula 7 in the article below does as well:

  1. Kolouri et al (2017) "Optimal Mass Transport: Signal-processing and machine learning applications" https://www.math.ucdavis.edu/~saito/data/acha.read.s19/kolouri-etal_optimal-mass-transport.pdf