ELBOing Stein: Variational Bayes with Stein Mixture Inference

Ola Rønning, Eric Nalisnick, Christophe Ley, Padhraic Smyth, Thomas Hamelryck

Keywords: variational inference, particle-based inference, variance collapse

TL;DR: We present Stein mixture inference, a particle-based inference method, that mitigates variance collapse in moderetly sized models.

Abstract:

Stein variational gradient descent (SVGD) (Liu & 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 & 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.

Primary Area: probabilistic methods (Bayesian methods, variational inference, sampling, UQ, etc.)

Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.

Submission Guidelines: I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2025/AuthorGuide.

Anonymous Url: I certify that there is no URL (e.g., github page) that could be used to find authors’ identity.

No Acknowledgement Section: I certify that there is no acknowledgement section in this submission for double blind review.

Submission Number: 6490