Gradient Flows for Sampling: Mean-Field Models, Gaussian Approximations and Affine Invariance

TL;DR

This work develops a unified framework for sampling distributions known up to normalization by evolving densities via gradient flows of the KL divergence under Fisher–Rao, Wasserstein, and Stein metrics. It introduces affine-invariant variants of these gradient flows, derives corresponding mean-field and Gaussian-approximate dynamics, and establishes convergence properties across several posterior regimes. Key contributions include the KL divergence's unique normalization-invariance property, the construction of affine-invariant metrics and flows, and the equivalence between Gaussian-approximate dynamics and moment-closure/proximal perspectives, all supported by numerical experiments. The results demonstrate that affine invariance enhances convergence for anisotropic targets and provide practical pathways for Gaussian variational inference and mean-field sampling in Bayesian inverse problems.

Abstract

Sampling a probability distribution with an unknown normalization constant is a fundamental problem in computational science and engineering. This task may be cast as an optimization problem over all probability measures, and an initial distribution can be evolved to the desired minimizer dynamically via gradient flows. Mean-field models, whose law is governed by the gradient flow in the space of probability measures, may also be identified; particle approximations of these mean-field models form the basis of algorithms. The gradient flow approach is also the basis of algorithms for variational inference, in which the optimization is performed over a parameterized family of probability distributions such as Gaussians, and the underlying gradient flow is restricted to the parameterized family. By choosing different energy functionals and metrics for the gradient flow, different algorithms with different convergence properties arise. In this paper, we concentrate on the Kullback-Leibler divergence after showing that, up to scaling, it has the unique property that the gradient flows resulting from this choice of energy do not depend on the normalization constant. For the metrics, we focus on variants of the Fisher-Rao, Wasserstein, and Stein metrics; we introduce the affine invariance property for gradient flows, and their corresponding mean-field models, determine whether a given metric leads to affine invariance, and modify it to make it affine invariant if it does not. We study the resulting gradient flows in both probability density space and Gaussian space. The flow in the Gaussian space may be understood as a Gaussian approximation of the flow. We demonstrate that the Gaussian approximation based on the metric and through moment closure coincide, establish connections between them, and study their long-time convergence properties showing the advantages of affine invariance.

Figures (8)

• Figure 1: Gaussian posterior case: convergence of different gradient flows in terms of the $L_2$ error of $\mathbb{E}[\theta_t]$, the relative Frobenius norm error of the covariance $\frac{\lVert \mathrm{Cov}[\theta_t] - \mathrm{Cov}[\theta_{{\rm true}}]\rVert_F}{\lVert \mathrm{Cov}[\theta_{\rm true}]\rVert_F}$, and the error of $\mathbb{E}[\cos(\omega^T \theta_t + b)]$.
• Figure 2: Gaussian posterior case: convergence of different dynamics in terms of $L_2$ error of $\mathbb{E}[\theta_t]$, the relative Frobenius norm error of the covariance $\frac{\lVert \mathrm{Cov}[\theta_t] - \mathrm{Cov}[\theta_{{\rm true}}]\rVert_F}{\lVert \mathrm{Cov}[\theta_{\rm true}]\rVert_F}$, and the error of $\mathbb{E}[\cos(\omega^T \theta_t + b)]$.
• Figure 3: Logconcave posterior case: convergence of different gradient flows in terms of the $L_2$ error of $\mathbb{E}[\theta_t]$, the relative Frobenius norm error of the covariance $\frac{\lVert \mathrm{Cov}[\theta_t] - \mathrm{Cov}[\theta_{{\rm true}}]\rVert_F}{\lVert \mathrm{Cov}[\theta_{\rm true}]\rVert_F}$, and the error of $\mathbb{E}[\cos(\omega^T \theta_t + b)]$.
• Figure 4: Logconcave posterior case: convergence of different dynamics in terms of $L_2$ error of $\mathbb{E}[\theta_t]$, the relative Frobenius norm error of the covariance $\frac{\lVert \mathrm{Cov}[\theta_t] - \mathrm{Cov}[\theta_{{\rm true}}]\rVert_F}{\lVert \mathrm{Cov}[\theta_{\rm true}]\rVert_F}$, and the error of $\mathbb{E}[\cos(\omega^T \theta_t + b)]$
• Figure 5: General posterior case: particles obtained by different gradient flows at $t=15$. Grey lines represent the contour of the true posterior.

Theorems & Definitions (74)

• Remark 2.1
• Theorem 2.2
• Remark 2.3
• Remark 3.1
• Definition 3.2
• Definition 3.3: Affine Invariant Gradient Flow
• Definition 3.4: Affine Invariant Metric
• Proposition 3.5
• Remark 3.6
• Remark 3.7