Mode Collapse of Mean-Field Variational Inference

TL;DR

This work explains why mean-field variational inference (MFVI) can exhibit mode collapse for bimodal targets by introducing ε-separatedness and proving MFVI concentrates mass on a single component when the components are sufficiently separated. To overcome this limitation, the authors propose Rotational Variational Inference (RoVI), which augments MFVI with an orthogonal rotation and jointly optimizes over the rotation and a product-measure, yielding tighter KL bounds for two-component Gaussian mixtures. They provide a gradient-based RoVI algorithm, with theoretical KL-bounds and practical evidence showing RoVI recovers multiple modes where MFVI fails. The results offer a principled path to mitigate mode collapse in non-log-concave VI problems and motivate further study of rotation- and transformation-based augmentations in variational inference.

Abstract

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. Specifically, when the target measure $π$ is a mixture $π= w P_0 + (1 - w) P_1$, 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. We introduce the notion to capture the separatedness of the two mixture components -- called $\varepsilon$-separateness -- and derive explicit bounds on the fraction of mass that any MFVI optimizer assigns to each component when $P_0$ and $P_1$ are $\varepsilon$-separated for sufficiently small $\varepsilon$. Our results suggest that the occurrence of mode collapse crucially depends on the relative position of the components. To address this issue, we propose the rotational variational inference (RoVI), which augments MFVI with a rotation matrix. The numerical studies support our theoretical findings and demonstrate the benefits of RoVI.

Figures (4)

• Figure 1: $\pi = \frac{1}{2}\,\mathcal{N}\left([-m,-m]^\top, I_2 \right) + \frac{1}{2}\,\mathcal{N}\left([m,m]^\top, I_2 \right)$. The above figure shows the coutour plots of $\pi$ and a MFVI optimizer when $m \in \{1,2,3\}$. We observe the mode-collapse is prominent when $m=3$. Moreover, when $m=3$, the MFVI optimizer can switch from one mode to the other depending on the initialization.
• Figure 2: Example of well-separated $P_0$ and $P_1$. The blue and red regions represent the supports of $P_0$ and $P_1$, respectively.
• Figure 3: Comparison of LMC, MFVI, and RoVI across six Gaussian mixture settings with varying means and covariances. RoVI yields high-quality approximations to the target distribution that closely match the marginal distributions of LMC samples, while MFVI collapses to a single mode under misalignment or strong anisotropy.
• Figure 4: The target is a symmetric Gaussian mixture: $\frac{1}{2}\,\mathcal{N}\left([-m,-m]^\top, I_2 \right) + \frac{1}{2}\,\mathcal{N}\left([m,m]^\top, I_2 \right)$. RoVI correctly recovers both components and the marginal distributions while MFVI collapses to a single mode.

Theorems & Definitions (20)

• Definition 2.1
• Lemma 2.2
• proof
• Theorem 2.3
• Remark 2.4
• Corollary 2.5
• Corollary 2.6
• Corollary 2.7
• Remark 2.8
• Proposition 3.1