ELBOing Stein: Variational Bayes with Stein Mixture Inference

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.

Thank you for highlighting this paper. We will add a discussion of ADVI with mixtures to the section Related Work in the article.

How does SMI compare with ADVI with mixtures?

In the article you highlighted, the authors propose two objective functions: SIELBO and SIWAE. SIELBO is similar to only using the "attractive force" of SMI (i.e., let the mixture weight alpha_k in [1] be the kernel evaluation and use the ELBO from line 273), which they point out suffers from particle collapse. They address this issue by introducing importance weighting, which leads to the SIWAE objective

In contrast, we tackle particle collapse by introducing a regularizer on the particles using the reproducing kernel. Our method can easily incorporate the SIWAE objective by adding importance weights to the ELBO of the attractive force. Since they have already demonstrated that this is an ELBO, it will be straightforward to show that, with our entropic regularizer, we still maintain the ELBO property.

The key difference between SIWAE and our method is that SIWAE's advancements conclude at this point. However, with SMI, we have a vast design space of kernels to explore, which may allow us to scale this method to larger models. The current article provides the necessary foundation for this further body of work (see below).

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.

Our experiments are designed to support our claim that SMI can mitigate variance collapse in SVGD. To this end, we are already considering models substantially larger than those used in previous work [2,3] on variance collapse in SVGD. By demonstrating SMI's robustness in moderately-sized models, we establish a solid foundation and necessary precondition for its extension to high-dimensional, real-world applications, which will require additional work on issues such as hyperparameter choice (notably regarding the choice of the kernel, which potentially provides a way to tackle the important open issue of the non-identifiability of deep models).

How is each component in the mixture distributions parameterized?

In a Gaussian mixture model (GMM), the overall distribution is represented as a combination of multiple Gaussian components, each of which is defined by specific parameters—namely, a mean and a variance. In SMI, each particle refers to a vector containing the mean and variance of one Gaussian component, effectively parameterizing it. Thus, a particle in this context is the pair (mean, variance) that defines the central location and spread of that Gaussian. When multiple particles are used, each representing a different Gaussian with unique parameters, they collectively form the GMM, with each particle contributing one component to the mixture. This setup allows the mixture model to flexibly approximate complex distributions by combining the influences of multiple Gaussian distributions, each defined by a particle.

What is the point of using Stein?

1. Optimizing a variational distribution that is a mixture model is inherently challenging, and NSVGD provides a principled and tractable framework for this task.
2. By leveraging NSVGD, SMI connects to a rich field of theoretical results, such as its relationship with Wasserstein flows, which provide a sound methodological foundation and allow for future research and innovation.
3. Mixture models are universal approximators of smooth densities. Consequently, SMI retains considerable freedom from strong parametric assumptions compared to methods relying on a single variational distribution, thereby preserving the flexibility of Stein-based methods.
4. The kernel component in SMI is a versatile tool that allows users to incorporate inductive biases into the variational approximation. For instance, the non-identifiability of deep models is a significant unresolved challenge in their Bayesian estimation [4]. By choosing suitable kernels, SMI offers the potential to mitigate such issues.

References

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.
2. Ba, Jimmy, et al. "Understanding the variance collapse of SVGD in high dimensions." International Conference on Learning Representations. 2021.
3. Zhuo, Jingwei, et al. "Message passing Stein variational gradient descent." International Conference on Machine Learning. PMLR, 2018.
4. Roy, Hrittik, et al. "Reparameterization invariance in approximate Bayesian inference." NeurIPS. 2024.

Could you provide additional analysis or intuition on why the combination of an…

To address variance collapse, the key distinction between SVGD and SMI lies in the space the particles occupy. SMI particles operate in a higher-dimensional space than SVGD particles. This allows the repulsive term ($\nabla_1 k(x,y)$) in SMI to influence both the shape of the distribution and its parameterized location. In contrast, SVGD particles can only control location.

In SVGD, each particle represents a latent parameter sample. Meanwhile, in SMI, each particle parameterizes an entire distribution. For example, if the parameterized distribution is a factorized Gaussian, each SMI particle would represent both the mean (location) and variance of the Gaussian. While the location component of an SMI particle shares the same space as an SVGD particle, the variance component has no equivalent in SVGD. As a result, the repulsive force in SMI operates in a broader space, encompassing both location and variance.

This distinction becomes evident when comparing SVGD and SMI in a two-particle approximation of a standard Gaussian distribution. The SMI approximation forms a Gaussian mixture. By breaking SMI particles into their location and variance components, we can visualize the location component within the same space as SVGD particles and the target Gaussian density. In this visualization, SMI particle locations converge toward the center of the target Gaussian, while SVGD particles spread out, maintaining equal distances from the Gaussian center.

Mean estimation image.

At first glance, it might seem that SMI particles have collapsed when considering only their locations. However, this interpretation is incomplete because it only tells half the story. When we examine the variance component of the SMI particles, we observe that a single SMI particle captures the variance in one dimension of the target Gaussian distribution, and both particles cover the variance in the other dimension. We will add this example to the appendix.

Variance estimation image.

Next, we explore the challenges posed by the RBF kernel in SMI, which exhibits certain pathological behaviors. To extend SMI to larger models, we must carefully design new kernels. Ideally, these kernels should account for the volume of the distribution parameterized by each SMI particle. We believe that probability product kernels, along with concepts from linear-Laplace approximations, will enable us to address these issues effectively.

Scaling SMI to large models, however, is a substantial research endeavor that requires dedicated study and goes beyond the scope of this work. This paper serves as foundational work, demonstrating that SMI can resolve variance collapse issues inherent to SVGD. The contributions here should be evaluated on this primary achievement.

Is there evidence that the entropy regularization term alone is insufficient to address variance collapse?

Yes. This corresponds to having the score function in the attractive force. With this substitution, we are doing SVGD (as shown in Appendix A3.1), which suffers from variance collapse, as demonstrated in the estimating Gaussian variance experiment and 1D regression with synthetic data.

SVGD with resampling (Ba et al. 2021)

We omitted SVGD with resampling (reSVGD), Algorithm 1, of Ba et al. (2021) in our Gaussian variance estimation experiment because in reproducing their results on the Gaussian variance estimation, we found that their proposed algorithm gives a biased estimator (i.e., $E[X \sim \text{reSVGD}] \not = 0]$). This is not the case for SMI, SVGD or ASVGD. We will include the reproduced result of their method for the Gaussian variance estimation and add the plot showing that reSVGD is biased in the appendix.

For the 1D and UCI experiments, we did not include a comparison with the Ba et al. (2021) algorithm due to several limitations. First, there are no publicly available implementations of their method, and the paper explicitly states that it is an analytic tool, not scalable due to its $O(m^4)$ computational complexity.

Second, the algorithm is inherently sequential, making it poorly suited for GPU acceleration, as it limits parallelization. Additionally, the memory overhead is substantial—each step requires sampling and storing a new set of particles, which can quickly lead to out-of-memory errors. Furthermore, frequent memory access on GPUs significantly diminishes performance gains.

Finally, their algorithm relies on knowing the variance to maintain during resampling, which is straightforward when the target is a standard Gaussian but impractical for the posterior of a Bayesian Neural Network (BNN). We did not include reSVGD in the UCI and 1D regression experiments for these reasons.

We would like to clarify that Figure 2 does not imply that variance collapse is mitigated with a small alpha ($0<\alpha\ll1$) for SMI in the general case.

In Figure 2 (right), we are using one particle, so the choice of alpha does not affect this experiment. We show this theoretically and experimentally. The theoretical validation is in Appendix A.3.2, where we show that one particle SMI is MFVI, so the objective SMI is the regular ELBO from eq. 4 (which contains no alpha). It is empirically validated in Figure 2 right, where the variance estimation is one regardless of $\alpha$.

In Figure 2 (left and middle), we investigate an artificial pathological case by using an overly large number of particles (20). That is, the variational model is much richer than warranted by the data. In this case, SMI will actually overestimate variance. However, we can still get good results for SMI by setting alpha to a small value (which mitigates the overestimation of variance in this pathological case). In the general case, this does not apply, and the entropy terms remain important.

HMC Baseline

Adding HMC as a baseline is an excellent suggestion. We will provide results on HMC with NUTS for the UCI datasets and 1D regression tasks. As expected, NUTS is comparatively slow but its prediction performance is good. We will include NUTS as a gold standard for comparison in the article and provide the results as an official comment.

W1: Variational inference offers a compelling alternative to sampling methods like MCMC due to its efficiency...

We appreciate the reviewer’s observation regarding the scope of our experiments. Our primary objective in this work is to introduce SMI and show that it effectively addresses the issue of variance collapse in SVGD. To this end, we focus on models that are already significantly larger than those typically considered in prior studies on variance collapse in SVGD (and that appeared in ICML and ICLR [1,2]). Our current contribution thus introduces a new method and demonstrates its robustness in addressing this specific challenge.

Scaling SMI to large-scale models is indeed important but extends beyond the scope of this paper. However, we argue that it is clear that SMI provides ample opportunities for scaling to larger models, notably through the use of bespoke kernels that go beyond simple RBF kernels, building inductive bias into the guide, hyperparameter tuning and so on. We plan to explore these opportunities in a future publication.

W2: The paper lacks theoretical insights or guidance on how SMI’s performance depends on the number of particles...

We do suggest using alpha to check if the approximation is too rich. This heuristic consists in checking if the variance (or calibration) of the estimator is invariant for values of alpha << 1. If the variance is stable, it suggests that too many particles are used, allowing inference with fewer particles without compromising uncertainty estimation. Beyond this heuristic, standard hyperparameter optimization techniques, such as cross-validation, dataset splitting, or leveraging information criteria, can also guide the selection of particle count.

We acknowledge, however, that developing a rigorous theoretical framework to determine the optimal number of particles remains an open and intriguing avenue for future research. Such work could include deriving bounds on variance or error as functions of particle count and alpha, which would offer deeper insights into this critical aspect of SMI. In this respect, the considerable body of theoretical results concerning NSVD and Stein’s method provides a solid ground. However, this is beyond the scope of this article, which introduces SMI and demonstrates that SMI can be used to mitigate variance collapse in SVGD.

Q1: It’s challenging to distinguish the lines representing different methods in Figure 2 (e.g. , )...

Thank you for the recommendation. We will update Figure 2 with distinct colors.

Q2: The experiments in Section 6.1 demonstrate that SMI overcomes variance collapse...

Our target posterior is a multivariate standard Gaussian, and the SMI approximation is modeled as a Gaussian Mixture Model (GMM) with mean-field Gaussian components. This setup is designed to align the posterior approximation with the true posterior. To further demonstrate this, we will include a plot in the appendix comparing the Frobenius norm between the estimated covariance matrices and the target covariance matrix. This will provide a quantitative measure of how well the SMI approximation captures the shape of the posterior.

Posterior shape plot.

References

1. Ba, Jimmy, et al. "Understanding the variance collapse of SVGD in high dimensions." International Conference on Learning Representations. 2021.
2. Zhuo, Jingwei, et al. "Message passing Stein variational gradient descent." International Conference on Machine Learning. PMLR, 2018.

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?

We are actively working on scaling SMI to larger models, and our current findings indicate that architectural depth, rather than the sheer number of parameters, presents a key challenge. We believe the problems with depth are related to the issue of non-identifiability [1], which complicates Bayesian inference considerably. We believe these challenges can be addressed by improvements in kernel design, which is beyond the scope of the current contribution.

We currently know that SMI begins to struggle with a BNN that has three or more layers, such as a fully connected 3-layered MLP with hidden dimension 100 (109,400 parameters) trained on MNIST. In contrast, a convolutional architecture like LeNet-1 (1995 version), which has only around 2,500 parameters but 8 layers, also proves challenging due to its depth. On the other hand, for a one-layer BNN, width does not appear to be a limiting factor for SMI's performance. These observations highlight the importance of addressing depth-related challenges to scale SMI to more complex architectures.

The paper lacks ablations to evaluate what happens as an underlying BNN gets deeper...

We appreciate the reviewer's suggestion regarding ablations to evaluate SMI's scalability to deeper BNNs. Scaling SMI to large models is indeed an important research direction. However, understanding how SMI scales to large models and more complex architectures is a substantial undertaking that goes beyond the scope of this work. This paper is intended as foundational research, focusing on demonstrating SMI's ability to address the critical issue of variance collapse in SVGD.

By resolving this core limitation, we establish a strong basis for further exploration, including scaling to deeper networks and incorporating last-layer BNNs. While we acknowledge the importance of such extensions, we believe that the contributions of this work should primarily be evaluated on the successful resolution of variance collapse and its implications for Bayesian inference quality. We view this work as a stepping stone, and future research will build upon these findings to address the challenges of scaling SMI to larger models.

Q2: What are the increased runtime costs compared to compared baselines?

To address this, we will provide per-step timings and commentary in the appendix for clarity and reproducibility. Below, we reproduce the per-step average inference time [sec/step] on the UCI datasets for a range of methods, including SMI. On UCI datasets, SMI exhibits slower inference compared to the VI-based baselines. A portion of this overhead arises from JIT compilation, which we believe can be reduced by optimizations in future releases of SMI.

When considering the recovery point experiment (Table 2), SMI demonstrates significantly improved runtime efficiency. On the mid-sized network, SMI achieves inference times 6x faster than SVGD. This observation suggests that while VI methods excel in runtime on UCI datasets, SMI provides a better trade-off when factoring in performance gains. Thus, in contexts where accuracy and robustness are critical, SMI is preferable despite its higher initial runtime cost.

We believe these observations underline the versatility of SMI, and we aim to address current runtime bottlenecks in future updates.

Citing Agarap (2018) in l478 as a reference for ReLUs seems rather odd...

Thank you for bringing this mistake to our attention. The reference was supposed to refer to the paper introducing ReLUs as an activation function in Deep Learning. As far as we know, this is the correct reference:

Fukushima, Kunihiko. "Neocognitron: A self-organizing neural network model for a mechanism of pattern recognition unaffected by shift in position." Biological cybernetics 36.4 (1980): 193-202.

References

1. Roy, Hrittik, et al. "Reparameterization invariance in approximate Bayesian inference." NeurIPS. 2024.

The experiments are rather small-scale and limited to regression data sets...

You are correct that our experiments are specifically designed to demonstrate how SMI can address variance collapse in SVGD. To achieve this, we use models that are significantly larger than those examined in prior studies on variance collapse in SVGD that appeared in ICML and ICLR [1,2]. The results in this study provide the necessary foundation for further work, notably on kernels and hyperparameter settings that make an application to larger problems possible.

Adding HMC as a baseline is an excellent suggestion. We will provide results on HMC with NUTS for the UCI datasets and 1D regression tasks. As expected, NUTS is comparatively slow, but its prediction performance is good. We will include NUTS as a gold standard for comparison in the article and provide the results as an official comment here. We will also include an application to classification (see below).

The experiments are limited to regressions with a homoscedastic, known observation noise.

While it is accurate that our experiments focus on homoscedastic models, the claim that we exclusively use known observation noise is incorrect. In our synthetic 1D regression experiments, the observation noise is indeed known. However, in the UCI experiments, we incorporate a prior over the noise, allowing each method to infer a posterior distribution for the noise parameter. Although the noise parameter is shared across all examples—ensuring the models remain homoscedastic—the key distinction is that the observation noise is not predetermined but rather inferred from the data.

What about classification or heteroscedastic regression tasks?

Heteroscedastic regression is commonly handled in BNNs by allowing the network to output the noise level as a function of the inputs. However, in our UCI experiments, we opted to adhere to the standard setup used in the UCI regression benchmark, which involves a homoscedastic noise model. Consequently, we did not employ a BNN architecture designed for heteroscedastic regression in this specific context.

To address the concern about classification, we will include results for a classification task using MNIST with a 1 and 2-layer BNN. This addition will demonstrate the applicability of our approach to classification tasks and provide an evaluation beyond regression.

l233 lacks a second closing bracket.

Thanks for bringing this to our attention. We have fixed the typo.

References

1. Ba, Jimmy, et al. "Understanding the variance collapse of SVGD in high dimensions." International Conference on Learning Representations. 2021.
2. Zhuo, Jingwei, et al. "Message passing Stein variational gradient descent." International Conference on Machine Learning. PMLR, 2018.