On Representations of Mean-Field Variational Inference

TL;DR

This paper introduces a comprehensive framework for analyzing mean-field variational inference (MFVI) algorithms, representing them through gradient flows, PDEs, and diffusion processes, and providing convergence guarantees.

Contribution

It develops a novel analytical framework for MFVI, connecting it with gradient flows and PDEs, and establishes convergence guarantees for practical algorithms.

Findings

MFVI can be represented as gradient flows on Wasserstein space.

Discretized coordinate ascent algorithms converge to the gradient flow.

Framework applies to various variational inference algorithms, ensuring convergence.

Abstract

The mean field variational inference (MFVI) formulation restricts the general Bayesian inference problem to the subspace of product measures. We present a framework to analyze MFVI algorithms, which is inspired by a similar development for general variational Bayesian formulations. Our approach enables the MFVI problem to be represented in three different manners: a gradient flow on Wasserstein space, a system of Fokker-Planck-like equations and a diffusion process. Rigorous guarantees are established to show that a time-discretized implementation of the coordinate ascent variational inference algorithm in the product Wasserstein space of measures yields a gradient flow in the limit. A similar result is obtained for their associated densities, with the limit being given by a quasi-linear partial differential equation. A popular class of practical algorithms falls in this framework, which provides tools to establish convergence. We hope this framework could be used to guarantee convergence of algorithms in a variety of approaches, old and new, to solve variational inference problems.