Projected Langevin dynamics and a gradient flow for entropic optimal transport

TL;DR

This paper introduces a diffusion process that samples from entropic optimal transport couplings, converging exponentially fast to the optimal solution and viewed as a Wasserstein gradient flow on the space of couplings.

Contribution

It develops a Langevin-type dynamics constrained to the set of couplings for entropic optimal transport, with proven exponential convergence and geometric interpretation.

Findings

The process remains in the set of couplings if initialized there.

It converges exponentially fast to the entropic optimal transport solution.

The dynamics can be viewed as a Wasserstein gradient flow on the space of couplings.

Abstract

The classical (overdamped) Langevin dynamics provide a natural algorithm for sampling from its invariant measure, which uniquely minimizes an energy functional over the space of probability measures, and which concentrates around the minimizer(s) of the associated potential when the noise parameter is small. We introduce analogous diffusion dynamics that sample from an entropy-regularized optimal transport, which uniquely minimizes the same energy functional but constrained to the set $\Pi(\mu,\nu)$ of couplings of two given marginal probability measures $\mu$ and $\nu$ on $\mathbb{R}^d$, and which concentrates around the optimal transport coupling(s) for small regularization parameter. More specifically, our process satisfies two key properties: First, the law of the solution at each time stays in $\Pi(\mu,\nu)$ if it is initialized there. Second, the long-time limit is the unique solution of an entropic optimal transport problem. In addition, we show by means of a new log-Sobolev-type inequality that the convergence holds exponentially fast, for sufficiently large regularization parameter and for a class of marginals which strictly includes all strongly log-concave measures. By studying the induced Wasserstein geometry of the submanifold $\Pi(\mu,\nu)$, we argue that the SDE can be viewed as a Wasserstein gradient flow on this space of couplings, at least when $d=1$, and we identify a conjectural gradient flow for $d \ge 2$. The main technical difficulties stems from the appearance of conditional expectation terms which serve to constrain the dynamics to $\Pi(\mu,\nu)$.