ELBOing Stein: Variational Bayes with Stein Mixture Inference

TL;DR

This paper introduces Stein Mixture Inference (SMI), a novel variational Bayesian inference method that extends SVGD by using mixture models to better estimate uncertainty, especially in small Bayesian neural networks.

Contribution

SMI generalizes SVGD by incorporating mixture models with user-specified guides, effectively avoiding variance collapse and requiring fewer particles for accurate uncertainty estimation.

Findings

SMI outperforms SVGD in avoiding variance collapse.

SMI requires fewer particles than SVGD for small BNNs.

SMI performs well on standard datasets.

Abstract

Stein variational gradient descent (SVGD) [Liu and Wang, 2016] performs approximate Bayesian inference by representing the posterior with a set of particles. However, SVGD suffers from variance collapse, i.e. poor predictions due to underestimating uncertainty [Ba et al., 2021], even for moderately-dimensional models such as small Bayesian neural networks (BNNs). To address this issue, we generalize SVGD by letting each particle parameterize a component distribution in a mixture model. Our method, Stein Mixture Inference (SMI), optimizes a lower bound to the evidence (ELBO) and introduces user-specified guides parameterized by particles. SMI extends the Nonlinear SVGD framework [Wang and Liu, 2019] to the case of variational Bayes. SMI effectively avoids variance collapse, judging by a previously described test developed for this purpose, and performs well on standard data sets. In addition, SMI requires considerably fewer particles than SVGD to accurately estimate uncertainty for small BNNs. The synergistic combination of NSVGD, ELBO optimization and user-specified guides establishes a promising approach towards variational Bayesian inference in the case of tall and wide data.