Annealing in variational inference mitigates mode collapse: A theoretical study on Gaussian mixtures
This addresses mode collapse in variational inference for practitioners, though it is incremental as it builds on existing annealing approaches.
The authors tackled mode collapse in variational inference by mathematically analyzing annealing strategies for Gaussian mixture models, deriving a formula for mode collapse probability and showing appropriate annealing can robustly prevent it, with numerical evidence extending findings to neural network models.
Mode collapse, the failure to capture one or more modes when targetting a multimodal distribution, is a central challenge in modern variational inference. In this work, we provide a mathematical analysis of annealing based strategies for mitigating mode collapse in a tractable setting: learning a Gaussian mixture, where mode collapse is known to arise. Leveraging a low dimensional summary statistics description, we precisely characterize the interplay between the initial temperature and the annealing rate, and derive a sharp formula for the probability of mode collapse. Our analysis shows that an appropriately chosen annealing scheme can robustly prevent mode collapse. Finally, we present numerical evidence that these theoretical tradeoffs qualitatively extend to neural network based models, RealNVP normalizing flows, providing guidance for designing annealing strategies mitigating mode collapse in practical variational inference pipelines.