Optimal transport and variational Bayesian inference

Abstract

We apply Lott–Sturm–Villani notion of Ricci curvature as well as McCann–Brenier theorem on optimal transport investigating mean-field variational inference (MFVI). In the continuous case where the probability rules are defined on Riemannian manifolds we obtain sufficient conditions for convexity of the corresponding variational problem. It can be shown that in the presence of symmetry the convexity can provide the preservation of mode. We then consider factorization into a mixed continuous-discrete case. By rescaling the Riemannian structure of the continuous component we have derived sufficient conditions under which the convexity of Kullback-Leibler functional is established. We apply the latter theorem to Bayesian Gaussian Mixture Model (GMM) and we introduce a new class of these models which have the advantage of being convex. It is shown that the MAP for the corrected GMM remains a consistent estimator. Additionally thanks to convexity we are able to prove the mode of MFVI for the new model in asymptotic regimes converges to the right solution. We do not know any other rigorous result for MFVI in statistical learning.