Mode Collapse of Mean-Field Variational Inference
This addresses a fundamental issue in variational inference for practitioners in machine learning and statistics, offering theoretical insights and a novel method to improve approximation accuracy in high-dimensional settings.
The paper tackles the problem of mode collapse in mean-field variational inference (MFVI) when approximating mixture distributions, providing the first theoretical explanation and showing that MFVI optimizers assign most mass to a single component depending on component separation. It proposes rotational variational inference (RoVI) to address this, with numerical studies demonstrating benefits.
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.