Bridging the Gap Between Variational Inference and Wasserstein Gradient Flows

TL;DR

This paper connects variational inference with Wasserstein gradient flows, showing how certain Wasserstein flows can be implemented as variational algorithms and introducing new gradient estimators for broader divergence measures.

Contribution

It demonstrates that Bures-Wasserstein gradient flows can be reformulated as Euclidean gradient flows, linking two optimization frameworks and extending gradient estimators to f-divergences.

Findings

Wasserstein gradient flows can be recast as variational inference algorithms.

Introduces a new gradient estimator for f-divergences.

Extends the framework to non-Gaussian variational families.

Abstract

Variational inference is a technique that approximates a target distribution by optimizing within the parameter space of variational families. On the other hand, Wasserstein gradient flows describe optimization within the space of probability measures where they do not necessarily admit a parametric density function. In this paper, we bridge the gap between these two methods. We demonstrate that, under certain conditions, the Bures-Wasserstein gradient flow can be recast as the Euclidean gradient flow where its forward Euler scheme is the standard black-box variational inference algorithm. Specifically, the vector field of the gradient flow is generated via the path-derivative gradient estimator. We also offer an alternative perspective on the path-derivative gradient, framing it as a distillation procedure to the Wasserstein gradient flow. Distillations can be extended to encompass $f$-divergences and non-Gaussian variational families. This extension yields a new gradient estimator for $f$-divergences, readily implementable using contemporary machine learning libraries like PyTorch or TensorFlow.