Mean-field variational inference (MFVI) is a widely used method for approximating high-dimensional probability distributions by product measures. It has been empirically observed that MFVI optimizers often suffer from mode collapse. For example, when the target measure π is a mixture π=wP0+(1−w)P1, the MFVI optimizer tends to place most of its mass near a single component of the mixture. This work provides the first theoretical explanation of mode collapse in MFVI. Starting from two-component mixtures, we introduce the notion of \emph{ε-separateness} to quantify the separation of the mixture components. We derive explicit bounds on the mass that any MFVI optimizer assigns to each component, when P0 and P1 are ε-separated for sufficiently small ε. We further extend this analysis to mixture target distributions with multiple components. Our results reveal that the occurrence of mode collapse crucially depends on the relative position of the components. To address this issue, we propose rotational variational inference (RoVI), which augments MFVI with a rotation matrix. The numerical studies support our theoretical findings and illustrate improved mode coverage with RoVI in Bayesian mixture posteriors, including a real-data example.
Figures & tables
Figure 1: π=21N([−m,−m]⊤,I2)+21N([m,m]⊤,I2) . The above figure shows the contour plots of π and a fitted MFVI approximation when m∈{1,2,3} . We observe the mode-collapse is prominent when m=3 . Moreover, when m=3 , the fitted approximation can switch from one mode to the other depending on the initialization.
Figure 2: Example of well-separated P0 and P1 . The blue and red regions represent the supports of P0 and P1 , respectively.
Figure 3: Comparison of LMC, MFVI, and RoVI across six Gaussian mixture settings with varying means and covariances. RoVI improves approximation in several settings where MFVI is underdispersed or collapses.
Figure 4: The target is a symmetric Gaussian mixture: 21N([−m,−m]⊤,I2)+21N([m,m]⊤,I2) . RoVI captures both components and the marginal distributions, while MFVI collapses to a single mode in the sufficiently separated cases.
Target
Method
m(q)
ℓ(q)
3D
MFVI
0.000±0.000
0.000±0.000
3D
RoVI
0.267±0.098
0.252±0.091
Old Faithful
MFVI
0.000±0.000
0.000±0.000
Old Faithful
RoVI
0.302±0.137
0.285±0.131
Table 1: Comparison of mode coverage for the simulated mixture (3D) and Old Faithful posterior. Reported values are means ± standard deviations across five paired runs following validation-based restart selection.
Figure 5: Joint densities for the Old Faithful posterior. We also report the probability P(μ1<μ2) . The columns show the reference and the fitted MFVI and RoVI approximations. Values of P(μ1<μ2) closer to 0.5 indicate more balanced coverage of the two symmetric modes.
Appendix figures & tables1 asset
Supplementary material from the paper’s appendix.
Appendix
Figure 6: Joint densities for the separated-mixture posterior (3D). The columns show the reference and the fitted MFVI and RoVI approximations, with the probability P(μ1<μ2) above each panel. RoVI retains both modes, while MFVI concentrates on one.
Mean Field Variational Inference (MFVI) is widely understood to underestimate posterior variance. By analysing conjugate Bayesian Linear Regression (BLR), we show that this characterization is incomplete: while MFVI underestimates the variance in parameter space, it can overestimate the predictive variance compared to the exact posterior. We show that if the MFVI posterior underestimates predictive variances in some directions, it necessarily overestimates them in others. Crucially, this overestimation occurs in directions where the training data concentrates. This leads to the surprising result that, for a test point drawn from the training distribution, MFVI's expected predictive variance exceeds that of the exact posterior. We demonstrate a pathological case of this effect, where the MFVI posterior fails to reduce predictive variance compared to the prior on in distribution data. We connect these results to the Cold Posterior Effect, arguing that varying the temperature can correct this overestimation, yielding predictions closer to those of the exact posterior. We validate our theory on synthetic and real-world regression tasks.
James Odgers, Ben Riegler, Siddharth Swaroop +1
University of Technology Nuremberg (UTN) · Helmholtz AI · Munich Center for Machine Learning (MCML) +2
In variational inference (VI), the practitioner approximates a high-dimensional distribution π with a simple surrogate one, often a (product) Gaussian distribution. However, in many cases of practical interest, Gaussian distributions might not capture the correct radial profile of π, resulting in poor coverage. In this work, we approach the VI problem from the perspective of optimizing over these radial profiles. Our algorithm radVI is a cheap, effective add-on to many existing VI schemes, such as Gaussian (mean-field) VI and Laplace approximation. We provide theoretical convergence guarantees for our algorithm, owing to recent developments in optimization over the Wasserstein space--the space of probability distributions endowed with the Wasserstein distance--and new regularity properties of radial transport maps in the style of Caffarelli (2000).
When approximating an intractable density via variational inference (VI) the variational family is typically chosen as a simple parametric family that very likely does not contain the target. This raises the question: Under which conditions can we recover characteristics of the target despite misspecification? In this work, we extend previous theoretical results on robust VI with location-scale families under target symmetries in two substantial ways: (1) We open them up to a wider range of divergences by providing sufficient conditions for exact recovery of the target mean and correlation matrix when using the forward Kullback-Leibler divergence and α-divergences. (2) By doing so, we find that we can drop the restrictive assumption of a log-concave target made in previous work, allowing us to give guarantees for a wider range of targets, including multi-modal ones. In our experiments, we show how our guarantees can serve as guidelines for the choice of the variational family and α-value and we illustrate on a diverse set of examples how and why optimization can fail in the absence of our sufficient conditions.