Bridging the Gap Between Variational Inference and Wasserstein Gradient Flows
This work addresses a theoretical gap for researchers in probabilistic machine learning, offering incremental insights by linking existing methods and providing practical tools.
The paper tackles the problem of connecting variational inference and Wasserstein gradient flows, showing that under certain conditions, the Bures-Wasserstein gradient flow can be reformulated as a Euclidean gradient flow, with its forward Euler scheme corresponding to the standard black-box variational inference algorithm, and extends this to derive a new gradient estimator for f-divergences.
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.