1D analytical solution for entropy-regularized Wasserstein distance

407 Views Asked by At

\begin{align} \mathcal{W}_{p}(\mu, \nu) &=\left(\int_{X} d^{p}\left(x, F_{\nu}^{-1}\left(F_{\mu}(x)\right)\right) \mathrm{d} \mu(x)\right)^{\frac{1}{p}} \\ &=\left(\int_{0}^{1} d^{p}\left(F_{\mu}^{-1}(z), F_{\nu}^{-1}(z)\right) \mathrm{d} z\right)^{\frac{1}{p}}\\ &=\left(\int_{0}^{1}\left|F_{\mu}^{-1}(z)-F_{\nu}^{-1}(z)\right|^{p} \mathrm{d} z\right)^{\frac{1}{p}} \end{align}

is the Kolouri (2017) closed-form analytical solution for the Wasserstein distance (optimal transport) when the source and target distributions are one-dimensional (1D)

(as opposed to 2D or higher dimensional, which would require linear programming)

Question

How can we similarly derive a closed-form solution for the entropy-regularized Wasserstein distance, also called Sinkhorn distance, for the 1D case?

\begin{equation} \mathcal{W}_\epsilon(\alpha, \beta) = \min_{\pi\in \Pi(\alpha\beta)} \int c(x,y) \mathrm{d}\pi(x,y) + \epsilon H(\pi) \end{equation} where $H(\pi) = \int \pi(x,y) \ln \pi(x,y)$ is the entropy of the transport plan $\pi(x,y)$.

Otherwise, how to prove that there can be no closed-form solution?

Notes

Breaking the problem into two parts, for example,

  • entropy of $\mathcal{W}_{p}(\mu, \nu)$, the unregularized Wasserstein distance, and
  • entropy of $H(\pi)$, the transport plan,

I am not aware of a closed-form solution for the second item.

Janati (2020) derives a closed-form solution for the formula in question in the case of two Gaussian distributions, which is fully differentiable unlike the unregularized case, but from what I can see in the article, no mention is made for a 1D solution.

1

There are 1 best solutions below

0
On

The best you can get in the general case is to write down a solution in terms of Lagrange multipliers- which will be functions, not scalars, in the continuous setting. The problem is equivalent to finding the maximum entropy joint distribution with equality constraints on the margins and an equality constraint on the expected cost. While we usually think about $\epsilon$ as part of an objective, it's functionally equivalent to an implicit constraint on the expected cost.

The solution will be in the form of

$$\pi(x,y) = \exp\left(\alpha(x)+\beta(y) - \frac{C(x,y)}{\epsilon}\right)$$

And the marginal constraints become a set of coupled integral equations. This is in essence a problem of solving a system of functional equations involving exponential family partition functions in full generality, which explains the difficulty. Outside of specific cases where there are known closed form joint distributions with the marginal constraints compatible with the integration of the cost (e.g. squared distance cost with Gaussian marginals giving a multivariate Gaussian or likewise for t or Laplace distributions) there is little hope for a nice closed form solution.