Transport map unadjusted Langevin algorithms: learning and discretizing perturbed samplers

TL;DR

This paper explores how transport maps can normalize target distributions to improve Langevin sampling, connecting continuous-time dynamics with Riemannian geometry and proposing new discretizations with convergence guarantees.

Contribution

It introduces a framework linking transport maps with Riemannian Langevin dynamics and irreversible perturbations, offering new methods for learning metrics and discretizations.

Findings

Transport maps induce Riemannian manifold Langevin dynamics.

Applying transport maps to ULA results in geometry-informed irreversible perturbations.

Discretized processes with transport maps can achieve convergence bounds in 2-Wasserstein distance.

Abstract

Langevin dynamics are widely used in sampling high-dimensional, non-Gaussian distributions whose densities are known up to a normalizing constant. In particular, there is strong interest in unadjusted Langevin algorithms (ULA), which directly discretize Langevin dynamics to estimate expectations over the target distribution. We study the use of transport maps that approximately normalize a target distribution as a way to precondition and accelerate the convergence of Langevin dynamics. We show that in continuous time, when a transport map is applied to Langevin dynamics, the result is a Riemannian manifold Langevin dynamics (RMLD) with metric defined by the transport map. We also show that applying a transport map to an irreversibly-perturbed ULA results in a geometry-informed irreversible perturbation (GiIrr) of the original dynamics. These connections suggest more systematic ways of learning metrics and perturbations, and also yield alternative discretizations of the RMLD described by the map, which we study. Under appropriate conditions, these discretized processes can be endowed with non-asymptotic bounds describing convergence to the target distribution in 2-Wasserstein distance. Illustrative numerical results complement our theoretical claims.