Convergence of drift-diffusion PDEs arising as Wasserstein gradient flows of convex functions

TL;DR

This paper establishes quantitative convergence rates for drift-diffusion PDEs as Wasserstein gradient flows of convex functions, revealing conditions under which these PDEs converge rapidly, with implications for optimization in probability measure spaces.

Contribution

It provides the first general convergence theory for these PDEs under convexity, including rates and applications to nonconvex problems and trajectory inference.

Findings

Suboptimality gap decreases as O(1/t) for convex objectives.

Faster than polynomial convergence for strongly convex objectives.

Applications demonstrated in entropy-regularized nonconvex optimization and trajectory inference.

Abstract

We study the quantitative convergence of drift-diffusion PDEs that arise as Wasserstein gradient flows of linearly convex functions over the space of probability measures on ${\mathbb R}^d$. In this setting, the objective is in general not displacement convex, so it is not clear a priori whether global convergence even holds. Still, our analysis reveals that diffusion {allows} a favorable interaction between Wasserstein geometry and linear convexity, leading to a general quantitative convergence theory, analogous to that of gradient flows in convex settings in the Euclidean space. Specifically, we prove that if the objective is convex and suitably coercive, the suboptimality gap decreases at a rate $O(1/t)$. This improves to a rate faster than any polynomial -- or even exponential in compact settings -- when the objective is strongly convex relative to the entropy. Our results extend the range of mean-field Langevin dynamics that enjoy quantitative convergence guarantees, and enable new applications to optimization over the space of probability measures. To illustrate this, we show quantitative convergence results for the minimization of entropy-regularized nonconvex problems, we propose and study an \emph{approximate Fisher Information} regularization covered by our setting, and we apply our results to an estimator for trajectory inference which involves the minimization of the relative entropy with respect to the Wiener measure in path space.