ELBOing Stein: Variational Bayes with Stein Mixture Inference

2024-10-30

arXiv.org

AI Summary

SVGD suffers from particle collapse in Bayesian inference, leading to underestimated posterior variance (“variance collapse”) and severely degrading uncertainty quantification—even for small- to medium-scale Bayesian neural networks (BNNs). To address this, we propose Stein Mixture Inference (SMI): each SVGD particle is modeled as a component of a learnable mixture distribution, jointly optimized via an ELBO lower bound regularized by a user-specified guiding distribution. SMI is the first framework to embed nonlinear SVGD into the variational Bayesian paradigm, synergistically mitigating collapse through particle-parameterized guidance and adaptive kernel design. Experiments on standard benchmarks demonstrate that SMI significantly improves uncertainty calibration. Notably, small BNNs achieve comparable or superior performance with far fewer particles than SVGD, while maintaining robustness in high dimensions and computational efficiency.

Technology Category

• Machine Learning: Calibration & Uncertainty Quantification
• Reasoning under Uncertainty: Probabilistic Programming
• Search and Optimization: Mixed Discrete/Continuous Search

Application Category

• Graph Algorithms and Modeling for the Web: Graph embeddings and representation learning for Web-related graphs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

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.

Problem

• Addresses variance collapse in SVGD for Bayesian inference
• Generalizes SVGD with mixture model components for uncertainty
• Optimizes ELBO with user-specified guides to improve predictions

Innovation

• Stein Mixture Inference with user-specified guides
• ELBO optimization for variational Bayesian inference
• NSVGD framework extension to prevent variance collapse