Gradient Flows for Sampling: Mean-Field Models, Gaussian Approximations and Affine Invariance (2302.11024v7)

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.

• The paper introduces a unified framework for sampling using gradient flows, leveraging the unique properties of KL divergence to address normalization challenges.
• The study employs affine-invariant metrics to ensure consistent convergence when sampling from highly anisotropic distributions.
• Gaussian approximations and mean-field dynamics are applied to derive closed-form gradient flows that enhance variational inference efficiency.

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

In this paper, the authors explore the application of gradient flows in the context of sampling from probability distributions that are specified up to a normalization constant. This problem is fundamental in computational science, finding applications in Bayesian inference and statistical mechanics, among other fields. The paper presents a comprehensive framework that unifies various sampling algorithms through the lens of gradient flows, focusing on different metrics like Fisher-Rao, Wasserstein, and Stein, while introducing the concept of affine invariance to enhance efficiency when dealing with highly anisotropic target distributions.

Key Contributions

1. Energy Functional Justification: The choice of Kullback–Leibler (KL) divergence as the energy functional is justified by its unique property among $f$-divergences: its first variation does not depend on the normalization constant. This is advantageous in scenarios where the normalization constant is unknown, such as Bayesian inference.
2. Affine Invariance: The paper introduces affine invariant metrics, which ensure that the convergence properties of the gradient flows remain unchanged under affine transformations of the parameter space. This is particularly beneficial for sampling from highly anisotropic distributions, where affine invariance can significantly improve convergence rates.
3. Gradient Flows in Probability Space: By employing Fisher-Rao, Wasserstein, and Stein metrics, the authors derive gradient flows that operate in the space of probability densities. The Fisher-Rao gradient flow, in particular, is noted for its diffeomorphism invariance, leading to a uniform exponential rate of convergence.
4. Gaussian Approximate Gradient Flows: Apart from operating in the full density space, the paper extends the notion of gradient flows to Gaussian approximations, which restrict variations to a family of Gaussian densities. This results in closed-form dynamics for the parameters defining the Gaussian densities, facilitating efficient variational inference methods.
5. Mean-Field Dynamics: The derived gradient flows correlate with mean-field dynamics, offering a potential bridge between gradient flow theory and particle-based sampling methods. Such dynamics are particularly relevant for computationally implementing these flows.

Implications and Future Directions

The introduction and rigorous treatment of affine invariance in gradient flows offer a promising avenue for designing more efficient sampling algorithms, especially in high-dimensional and poorly scaled problems. The work not only consolidates existing methodologies under a unified framework but also sets the stage for further exploration into other potential invariance properties that could be exploited in the development of sampling methods.

For future developments, extending the methodology to handle non-Gaussian posteriors more effectively remains a critical challenge. This includes addressing scenarios with complex posterior landscapes, such as multimodal distributions or distributions concentrating on manifolds with significant curvature. Additionally, integrating these methodologies with scalable inference frameworks, akin to those used in ensemble Kalman methods, can offer significant advantages in large-scale Bayesian inverse problems.

In summary, the paper significantly contributes to the literature on probabilistic inference by providing a deep theoretical foundation and practical approach for leveraging gradient flows in sampling tasks, with the potential to impact a variety of applications across computational science and machine learning fields.