ELBOing Stein: Variational Bayes with Stein Mixture Inference

Authors: Ola Rønning, Eric Nalisnick, Christophe Ley, Padhraic Smyth, Thomas Hamelryck
Primary area: probabilistic methods (Bayesian methods, variational inference, sampling, UQ, etc.)
Keywords: variational Bayes, particle-based inference, mixture 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.

Decision: accept (poster)
Reviewer scores: 8, 6, 6, 6

Reviews and discussion:
<thread>
creation date:
1730067600502

thread_type:
official_review

<thread_object>
creator: reviewer_dciu

<content>
summary:
The posterior of a Bayesian model is sometimes intractable, calling for approximate inference techniques. This paper focuses on the idea of approximating the Bayesian posterior with a mixture distribution where each component is parameterized separately but still in the same family. By viewing the ELBO as an objective with permutation invariant parameters, this paper incorporates ideas from Nonlinear-SVGD (NSVGD) and develop a Stein-style update rule of the mixture parameters. The resulted method, called Stein mixture inference (SMI), prevents itself from variance collapse. The paper also shows that asymptotically the optimized bound is an ELBO.

In the experiments, it is first shown that Stein-based methods suffer from variance collapse in synthesized Gaussian models and 1D regression models. In contrast, SMI and vanilla VI (OVI) produce the desired uncertainty. On the UCI regression benchmarks, SMI gives the best NLL in the most cases, comparing with other Stein-based methods and OVI.

strengths:
I believe the method is novel and the main idea is sound. The claims are clearly presented and supported by empirical evidences. Overall it is a complete work with a few concerns.

weaknesses:
This paper has the title of ELBOing Stein, but I would rather call it Steining ELBO, which seems a bit unnecessary. To convince me to increase my score, I would like to see a discussion of [1], which also uses a mixture distribution to approximate the posterior. 

If one is using a mixture distribution to target the Bayesian posterior, the most direct approach would be to try VI, instead of deriving a complicated Stein-based method. One pro of VI is that the mixture weights can be adjusted while one con is that it does not fully make use of the exchangeability of parameters. However, this work only considers mean-field VI, which is a really weak baseline. I would like to see how the permutation invariance helps the optimization of the mixture distribution.

It is not surprising that optimizing an VI objective prevents the approximate posterior from the pitfalls of SVGD. As shown in the paper, OVI does not have the issue. The argument that "SMI is more particle-efficient than SVGD" is translated to me as "VI with mixtures is more particle-efficient than SVGD". Then what is the point of using Stein?

Line 150 says that "Particle methods are attractive due to their freedom from strong parametric assumptions". Mixture distribution in this paper seems to be a strong parametric assumption, especially when it uses fewer particles than SVGD, which further drags this work away from Stein. 

The experiment section is also rather weak. The benchmark models all have very low dimensions. I expect a Bayesian inference algorithm in 2024 to be tested on more recent benchmarks, like models in posteriorDB, or larger BNN problems.

[1] Morningstar, W., Vikram, S., Ham, C., Gallagher, A., & Dillon, J. (2021, March). Automatic differentiation variational inference with mixtures. In International Conference on Artificial Intelligence and Statistics (pp. 3250-3258). PMLR.

questions:
- How does SMI compare with ADVI with mixtures? 
- How is each component in the mixture distributions parameterized?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730547556255

thread_type:
official_review

<thread_object>
creator: reviewer_hpkw

<content>
summary:
- The authors propose a method called Stein Mixture Inference (SMI) to address the issue of variance collapse observed in Stein Variational Gradient Descent (SVGD).
- SMI extends the Nonlinear SVGD framework (Wang & Liu, 2019) to variational inference (VI) by allowing each particle to parameterize a component distribution within a mixture model, thereby introducing ELBO-like objectives.
- The authors show that SMI offers several advantages over standard SVGD by effectively mitigating variance collapse.

strengths:
- The application of variational inference (VI) concepts to Stein Variational Gradient Descent (SVGD) appears novel and intriguing.
- The authors validate their VI-based approach through numerical experiments on several UCI benchmark datasets, demonstrating good performance. The results seem to suggest that this approach effectively mitigates the impact of variance collapse.

weaknesses:
### Insufficient Analysis of the Motivation Behind Extending SVGD with VI for Variance Collapse Mitigation:
- The main objective of this paper, as I understand it, is to mitigate variance collapse by extending the SVGD objective function through a combination of an ELBO-like objective from VI and the Non-linear SVGD framework. However, it is not entirely clear "why" this extension effectively mitigates variance collapse. While Figure 1 provides a conceptual illustration, it does not intuitively explain how the proposed method addresses variance collapse. Additionally, while the third paragraph of the Introduction discusses the motivation for this approach, it remains unclear how using a mixture of approximate distributions around SVGD particles with a VI-inspired objective avoids variance collapse.
- The authors propose controlling particles through variational distributions, similar to VI, as a solution to the variance collapse issue. However, given the use of the NSVGD framework, the critical role of this aspect remains unclear. The entropy regularization term could potentially affect not only mode collapse but also variance collapse. If the VI-inspired approach is indeed effective, the method should perform well even with $\alpha=0$. In this context, Figure 2 shows that variance collapse is mitigated even when $\alpha$ takes small values, suggesting that particle control via variational distributions may be effective. On the other hand, this result implies that regularization may not play a significant role, raising questions about the necessity of combining it with NSVGD. Overall, it remains unclear why the NSVGD framework is essential and which part of the proposed approach effectively addresses variance collapse.

### Concerns Regarding the Limited Number of Comparative Methods:
- For sample approximations of the posterior distribution, methods such as HMC (Neal, 2011) and MMD descent (Arbel et al., 2019; Ba et al., 2021) are also effective. However, this study only compares performance within the SVGD family of methods and EVI, leaving questions about the extent to which the proposed method mitigates variance collapse in the broader context of approximate inference. Given that (Ba et al., 2021) also includes these methods in numerical experiments addressing variance collapse, this comparison is essential for validating contributions in this research area.
- Additionally, the absence of a comparison with the resampling method proposed by (Ba et al., 2021) raises concerns regarding the integrity of the performance evaluation. While the authors argue in Section 5 that the resampling method is computationally infeasible, I believe this does not fully justify its exclusion as a comparative method. Given the availability of an “NVIDIA Quadro RTX 6000 GPU,” running such methods may not be computationally prohibitive, at least for datasets like the UCI benchmarks.
- Furthermore, I find it difficult to agree with the authors’ claim: “Annealing SVGD (ASVGD) D’Angelo & Fortuin (2021a) is the only alternative that directly addresses variance collapse in SVGD with a viable method.” I believe that the resampling method proposed by (Ba et al., 2021) is also aimed at mitigating the variance collapse problem.

### Citation:
- (Neal, 2011): R. M. Neal. MCMC using Hamiltonian dynamics. Handbook of Markov Chain Monte Carlo, 2(11):2. https://arxiv.org/abs/1206.1901.
- (Arbel et al., 2019): M. Arbel, A. Korba, A. Salim, and Arthur Gretton. Maximum Mean Discrepancy Gradient Flow. NeurIPS2019. https://arxiv.org/abs/1906.04370.

questions:
- Could you provide additional analysis or intuition on why the combination of an ELBO-like objective from VI and the Non-linear SVGD framework effectively mitigates variance collapse? Specifically, how does modeling a mixture of approximate distributions around the vicinity of SVGD particles help in avoiding variance collapse? A more detailed explanation or visual representation would be appreciated.
  - For example, could you provide a more detailed explanation or visual representation of how the mixture components interact with the ELBO objective to mitigate variance collapse? For instance, a step-by-step explanation or diagram illustrating how the proposed method addresses the variance collapse problem would be helpful.
- Why is the integration with the NSVGD framework necessary in your method? Is there evidence that the entropy regularization term alone is insufficient to address variance collapse? Given that Figure 2 shows variance collapse is mitigated even when $\alpha$ takes small values, does this imply that the regularization component may not be as critical? If so, what is the rationale for including it in the framework?
- Why were methods such as HMC and MMD descent not included in the comparative analysis, especially given their relevance in approximate inference and their use in experiments in (Ba et al., 2021)?
  - If possible, could you add comparisons with HMC (Neal, 2011) and MMD descent (Arbel et al., 2019; Ba et al., 2021) in the experimental section, particularly at least on the UCI datasets, to provide a broader context for evaluating SMI’s performance in addressing variance collapse? If a full comparison is not feasible, could you discuss how SMI might be expected to compare to these methods theoretically or empirically, based on existing literature?
- Could you elaborate on why the resampling method from (Ba et al., 2021) was excluded as a comparative method, despite the computational resources available (e.g., “NVIDIA Quadro RTX 6000 GPU”)? Is this method genuinely computationally infeasible for UCI benchmark datasets, or were there other factors influencing its exclusion?
  - So, could you include the resampling method from Ba et al. (2021) in your comparisons, particularly on the UCI datasets, to strengthen the evaluation? If this is not feasible, could you provide a more detailed justification for why it is computationally infeasible, even with the available GPU resources? Additionally, if an empirical comparison is truly not possible, could you discuss how SMI theoretically compares to the resampling approach?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730609344164

thread_type:
official_review

<thread_object>
creator: reviewer_mftt

<content>
summary:
The paper introduces Stein Mixture Inference (SMI), which optimizes a lower bound to the Evidence Lower Bound (ELBO). 
SMI extends Nonlinear Stein Variational Gradient Descent (NSVGD) to the variational Bayes setting and addresses the issue of variance collapse. 
The effectiveness of SMI is demonstrated on both synthetic and real-world datasets.

strengths:
1. The problem addressed in this paper is both important and compelling. Traditional approaches like ordinary mean-field variational inference (OVI) and Stein Variational Gradient Descent (SVGD) often experience variance collapse, whereas SMI provides more accurate variance estimates, improving uncertainty quantification.

2. The paper is well-written, providing a clear background and a thorough summary of related work. As someone slightly unfamiliar with the field, I particularly appreciated the authors' effort to re-explain and contextualize prior results, which greatly helped in assessing the paper's contributions.

3. SMI is compared with other methods across a variety of synthetic and real-world datasets.

weaknesses:
1. Variational inference offers a compelling alternative to sampling methods like MCMC due to its efficiency, especially in high-dimensional settings and with large-scale datasets. 
However, the current validation of SMI is limited to small to moderately-sized models, which somewhat limits its appeal and persuasiveness for broader, large-scale applications.

2. The paper lacks theoretical insights or guidance on how SMI’s performance depends on the number of particles $m$.
Providing recommendations or analysis on selecting an appropriate particle count $m$ would greatly enhance its practical applicability.

questions:
1. It’s challenging to distinguish the lines representing different methods in Figure 2 (e.g. $SMI_{1}$, $SMI_{20}$). 
Using distinct colors for each method would improve the visualization and make the differences clearer.
2. The experiments in Section 6.1 demonstrate that SMI overcomes variance collapse. It would also be valuable to assess whether the approximate distribution given by SMI accurately captures the shape of the posterior. 
This could be evaluated by comparing the estimated covariance matrix with the target covariance matrix.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730885902784

thread_type:
official_review

<thread_object>
creator: reviewer_ryzm

<content>
summary:
The authors improve Stein variational gradient descent (Liu & Wang, 2016) by extending nonlinear SVGD (Wang & Liu, 2019) by learning a density-based mixture model to approximate the posterior, instead of solely relying on particles (i.e., delta-distributions).
They evaluate their method on a set of (small) scale regression tasks.

strengths:
- The method is a straightforward and effective extension of SVGD/NSVGD
- The paper is well-written and easy to follow and the same goes for the provided codebase

weaknesses:
- The experiments are rather small-scale and limited to regression data sets. Their aim seems to be primarily to demonstrate the relative performance of the proposed approach compared to prior SVGD-related approaches rather than, its absolute performance. In the list of baselines, at least a comparison against an HMC performance on the UCI data sets would have been nice to see how close it can come to it (or improve upon it).
- The paper lacks ablations to evaluate what happens as an underlying BNN gets deeper, i.e., to what extent it can handle the increase in parameters. A deep experiment could be a combination with last-layer BNNs, i.e., learn the mixture not for the whole net, but treat only the penultimate layer in a Bayesian fashion.
- ~~The experiments are limited to regressions with a homoscedastic, known observation noise.~~ What about classification or heteroscedastic regression tasks? _Edit: The claimed fixed noise was an error on my side. See the answer of the authors on this_
- Citing Agarap (2018) in l478 as a reference for ReLUs seems rather odd. In their work, they evaluate the usage of a ReLU in place of a softmax for classification, i.e., nothing related to the current work nor has the ReLU been introduced in that paper.

### Typos
- l233 lacks a second closing bracket

questions:
- Q1: Can the authors speculate on performance as a function of the parameter count, e.g., sticking to BNNs, at which depth/width would the method start to struggle?
- Q2: What are the increased runtime costs compared to compared baselines?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732198238445

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
MNIST results

comment:
The tables below summarize the performance of 1-layer and 2-layer Bayesian Neural Networks (BNNs) on the MNIST dataset, evaluated across several metrics: confidence (Conf), negative log-likelihood (NLL), accuracy (Acc), Brier score (Brier), expected calibration error (ECE), and maximum calibration error (MCE).

For the 1-layer BNN, SMI outperforms other methods across all metrics except ECE and MCE. Given the robustness of the Brier score compared to ECE and MCE—which are sensitive to the number of bins (100 bins were used in this evaluation)—SMI is considered the best-calibrated method. When evaluating all metrics collectively, SMI stands out as the preferred approach.

For the 2-layer BNN, SMI again outperforms other methods on most metrics, except for the Brier score, which is on par with MAP. We regarded SMI and MAP as the best-calibrated methods for the same reasons outlined earlier. Overall, SMI remains the preferred approach when considering all metrics together.

We exclude HMC with NUTS from the MNIST analysis because HMC does not support subsampling, rendering inference computationally infeasible given our, or any conventional, hardware constraints.

## 1-layer BNN table

| Method | Conf ($\uparrow$)  	| NLL ($\downarrow$) 	| Acc ($\uparrow$)   	| Brier ($\downarrow$)   | ECE ($\downarrow$) 	| MCE ($\downarrow$) 	|
|:------:|:----------------------:|:----------------------:|:----------------------:|:----------------------:|:----------------------:|:----------------------:|
| ASVGD  | $0.972 \pm 0.002$  	| $0.053 \pm 0.004$  	| $0.949 \pm 0.003$  	| $0.074 \pm 0.005$  	| $0.135 \pm 0.007$  	| $0.634 \pm 0.024$  	|
| MAP	| $0.973 \pm 0.001$  	| $0.050 \pm 0.002$  	| $0.952 \pm 0.001$  	| $0.068 \pm 0.000$  	| $0.133 \pm 0.000$  	| $\bf{0.574 \pm 0.000}$ |
| OVI	| $0.921 \pm 0.006$  	| $0.158 \pm 0.012$  	| $0.908 \pm 0.006$  	| $0.106 \pm 0.007$  	| $\bf{0.085 \pm 0.010}$ | $0.630 \pm 0.136$  	|
| SMI	| $\bf{0.979 \pm 0.001}$ | $\bf{0.039 \pm 0.003}$ | $\bf{0.957 \pm 0.003}$ | $\bf{0.065 \pm 0.005}$ | $0.148 \pm 0.012$  	| $0.631 \pm 0.047$  	|
| SVGD   | $0.972 \pm 0.003$  	| $0.054 \pm 0.006$  	| $0.949 \pm 0.004$  	| $0.074 \pm 0.007$  	| $0.139 \pm 0.014$  	| $0.653 \pm 0.048$  	|

## 2-layer BNN table
Method | Conf ($\uparrow$)  	| NLL ($\downarrow$) 	| Acc ($\uparrow$)   	| Brier ($\downarrow$)   | ECE ($\downarrow$) 	| MCE ($\downarrow$)	 
:------:|:----------------------:|:----------------------:|:----------------------:|:----------------------:|:----------------------:|:----------------------:
 ASVGD  | $0.956 \pm 0.004$  	| $0.104 \pm 0.011$  	| $0.936 \pm 0.004$  	| $0.083 \pm 0.003$  	| $0.132 \pm 0.011$  	| $0.651 \pm 0.075$ 	 
 MAP	| $0.976 \pm 0.001$ | $0.044 \pm 0.003$ | $0.955 \pm 0.001$ | $\bf{0.066 \pm 0.000}$ | $0.126 \pm 0.000$ | $\bf{0.614 \pm 0.000}$
 OVI	| $0.913 \pm 0.005$  	| $0.182 \pm 0.012$  	| $0.899 \pm 0.005$  	| $0.116 \pm 0.005$ | $\bf{0.084 \pm 0.009}$ | $0.652 \pm 0.133$
 SMI	| $\bf{0.979 \pm 0.002}$ | $\bf{0.042 \pm 0.005}$ | $\bf{0.956 \pm 0.002}$ | $\bf{0.067 \pm 0.003}$ | $0.150 \pm 0.014$  	| $0.653 \pm 0.057$ 	 
 SVGD   | $0.960 \pm 0.004$  	| $0.091 \pm 0.011$  	| $0.940 \pm 0.002$  	| $0.081 \pm 0.004$  	| $0.135 \pm 0.013$  	| $0.649 \pm 0.044$
</content>
</thread_object>

</thread>
<thread>
creation date:
1732717887370

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Rebuttal Revisions

comment:
The reviewers' feedback and suggestions have directly contributed to improving the evaluation of our method as well as enhancing the overall clarity and quality of the work. Your comments and critiques have been valuable in refining both our methodology and presentation. Thank you. 

We have updated the article with the following:
- Added discussion of ADVI for mixture models to Related Works
- Added experiments showing that the Ba et al. (2021) sampling algorithm is biased to the appendix.
- Added HMC with NUTS results to UCI benchmark and 1D regression.
- Added MNIST results to the main article.
- Added forces discussion to the Appendix.
- Added UCI timings to the Appendix.
- Moved recovery point to the Appendix to make space for the MNIST experiment.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This study addresses the variance collapse issue in Stein variational gradient descent (SVGD) by proposing a novel method called Stein mixture inference (SMI), which combines existing Nonlinear SVGD with mixture models. The key property of this method is its ability to retain the nonparametric sample-based characteristic of SVGD, while simultaneously maximizing the evidence lower bound (ELBO), as commonly done in variational inference. This connection allows the proposed method to be interpreted within the framework of variational inference. However, as noted by multiple reviewers, the experiments presented in the paper are limited, and it is unclear how the maximization of ELBO contributes to resolving the variance collapse issue. Despite these weaknesses, the idea of linking SVGD with ELBO maximization is innovative and valuable to the community. With appropriate revisions addressing these weaknesses, the study would be considered suitable for acceptance.