Sampling via gradient flows in the space of probability measures

Sampling a target probability distribution with an unknown normalization constant is a fundamental challenge in computational science and engineering. Recent work shows that algorithms derived by considering gradient flows in the space of probability measures open up new avenues for algorithm developments. This paper makes three contributions to this approach to sampling, by scrutinizing the design components of such gradient flows. Any instantiation of a gradient flow for sampling needs an energy functional and a metric to determine the flow, as well as numerical approximations of the flow to derive algorithms. Our first contribution is to show that the Kullback-Leibler (KL) divergence, as an energy functional, has the <italic>unique</italic> property (among all <inline-formula content-type="math/mathml"> <mml:math xmlns:mml="http://www.w3.org/1998/Math/MathML" alttext="f"> <mml:semantics> <mml:mi>f</mml:mi> <mml:annotation encoding="application/x-tex">f</mml:annotation> </mml:semantics> </mml:math> </inline-formula> -divergences) that gradient flows resulting from it do not depend on the normalization constant of the target distribution; this justifies the widespread use of the KL divergence in sampling. Our second contribution is to study the choice of metric from the perspective of invariance. The Fisher-Rao metric is known as the <italic>unique</italic> choice (up to scaling) that is diffeomorphism invariant. As a computationally tractable alternative, we introduce a relaxed, affine invariance property for the metrics and gradient flows. In particular, we construct various affine invariant Wasserstein and Stein gradient flows. Affine invariant gradient flows are shown to behave more favorably than their non-affine-invariant counterparts when sampling highly anisotropic distributions, in theory and by using particle methods. Our third contribution is to study, and develop efficient algorithms based on Gaussian approximations of the gradient flows; this leads to an alternative to particle methods. We establish connections between various Gaussian approximate gradient flows, discuss their relation to gradient methods arising from parametric variational inference, and study their convergence properties. Our theory and numerical experiments demonstrate the strengths and potential limitations of the Gaussian approximate Fisher-Rao gradient flow, which is affine invariant, by considering a wide range of target distributions.