Computing Optimal Transport Plans via Min-Max Gradient Flows

TL;DR

The Kantorovich optimal transport problem is posed as a min-max problem with a Nash equilibrium that can be obtained dynamically via a two-player game, providing a framework for approximating optimal couplings and implementing the gradient descent algorithm with a particle method.

Abstract

We pose the Kantorovich optimal transport problem as a min-max problem with a Nash equilibrium that can be obtained dynamically via a two-player game, providing a framework for approximating optimal couplings. We prove convergence of the timescale-separated gradient descent dynamics to the optimal transport plan, and implement the gradient descent algorithm with a particle method, where the marginal constraints are enforced weakly using the KL divergence, automatically selecting a dynamical adaptation of the regularizer. The numerical results highlight the different advantages of using the standard Kullback-Leibler (KL) divergence versus the reverse KL divergence with this approach, opening the door for new methodologies.

Figures (4)

• Figure 1: The transport map generated with a kernel density estimator method from liu_approximating_2021 (top) is limited by the number of points ($N=1,000$), whereas estimating the density by counting voxels (bottom, our method) allows for a greater number of points ($20,000$) and faster computation time.
• Figure 2: In the setting of two log-concave input measures, the coupling cost starts small due to the initialization of $X_i\approx Y_i$, and increases to approach the optimal coupling cost. The regularization penalty weight $\Lambda$ increases more slowly over time, as expected due to the $\mathcal{O} (\sqrt{t})$ estimate, and the KL divergence between the marginals and the inputs decreases as the particles drift to minimize the entropy regularization terms.
• Figure 3: The source and target distributions for the experiment (b) are shown. Note that the input marginals are not log-concave, going beyond Assumption \ref{['assump:c_mu_nu_convex']}.
• Figure 4: We test our method on two input measures which are non-log-concave (Left side is $\mu$ and right side is $\nu$). When running our method with two different entropy regularizers, using $\text{KL}(\rho\,|\,\mu,\nu)$ for the $Y^{(2)}$ particles and $\text{KL}(\mu,\nu\,|\,\rho)$ for the $X^{(2)}$ particles (labeled "both KL") is compared with using the same for all particles. While the displacement convexity of $\text{KL}(\cdot\,|\,\mu)$ is well known, the displacement convexity of $\text{KL}(\mu\,|\,\cdot)$ is not.

Theorems & Definitions (29)

• Definition 2.1: Displacement Convexity
• Lemma 3.1
• proof
• Theorem 3.2: Displacement Convexity
• proof
• Theorem 3.3: Convergence
• proof
• Lemma 3.4: Marginals of best response
• proof
• Lemma 3.5: Convergence of Second Moments