ELBOing Stein: Variational Bayes with Stein Mixture Inference
This addresses a specific bottleneck in variational Bayesian inference for practitioners dealing with uncertainty estimation in models like small Bayesian neural networks, though it appears incremental as it builds on existing frameworks.
The paper tackled variance collapse in Stein variational gradient descent (SVGD) for Bayesian inference by introducing Stein Mixture Inference (SMI), which uses particles to parameterize mixture distributions and optimizes an ELBO, resulting in effective avoidance of variance collapse and requiring fewer particles than SVGD for small Bayesian neural networks.
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.