Research map / Probabilistic and scientific modeling / Variational and generative sampling

Annealing in variational inference mitigates mode collapse: a theoretical study on Gaussian mixtures

NeurIPS 2026 · Probabilistic and scientific modeling › Variational and generative sampling · approach: Amortized and online variational inference · 0 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Analyzes annealed reverse-KL variational inference on a bimodal isotropic Gaussian mixture with learnable variances, reducing gradient flow to a 3-variable ODE to derive a closed-form mode-collapse probability and validating the initial-temperature/annealing-time trade-off on RealNVP flows.

Problem

Mode collapse in variational inference over multimodal targets, and how annealing (tempering the target) should be scheduled to avoid it; little theory exists on annealing behavior.

Approach

Extends a summary-statistics analysis of VI on bimodal Gaussian mixtures to learnable variances via a quasi-static assumption (variance equals inverse temperature), yielding a low-dimensional ODE whose high- and low-temperature regimes give a closed-form collapse probability as a function of initial temperature and annealing time.

Key results

• Derived closed-form mode collapse probability p(beta) via integral over I(beta)=int sqrt(R^2 beta(s)/2pi) ds, matching experiments with alpha approx 0.608.
• For high initial temperature, collapse probability depends only on annealing rate, I ~ sqrt(2/pi) t0 / log(1/beta_i).
• RealNVP (d=128, 8 coupling layers) shows the same trade-off: increasing initial temperature requires increasing annealing time to avoid collapse.

Surprising findings

• Trainable variances alone do not avoid mode collapse, and a covering initialization with large variance still collapses because variances rapidly converge to 1 (target value).

Limitations

• Theory relies on quasi-static variance approximation sigma^2 = 1/beta, justified heuristically.
• Restricted to two-mode target with fixed equal weights and spherical constraint on means.
• Collapse threshold for RealNVP is an empirical variance-based metric; alpha fitted rather than derived.

Experimental setup

Datasets: synthetic bimodal Gaussian mixture target (d=512 and d=128)

Metrics: mode collapse probability variance of projected samples

Baselines: vanilla VI without annealing covering initialization annealing schedule beta(t)=1-(1-beta_i)(1-exp(-t/t0))

Builds on: JKO scheme Wasserstein geometry for Gaussian mixtures summary statistics analysis of VI for Gaussian mixtures RealNVP

Closest papers by meaning

• End-to-end Learning of Gaussian Mixture Priors for Diffusion Sampler (ICLR 2025 · similarity 0.90)
• Variational Inference with Mixtures of Isotropic Gaussians (NeurIPS 2025 · similarity 0.87)
• Efficient Mixture Learning in Black-Box Variational Inference (ICML 2024 · similarity 0.87)
• Theoretical Guarantees for Variational Inference with Fixed-Variance Mixture of Gaussians (ICML 2024 · similarity 0.86)
• Annealing Flow Generative Models Towards Sampling High-Dimensional and Multi-Modal Distributions (ICML 2025 · similarity 0.86)
• Particle Semi-Implicit Variational Inference (NeurIPS 2024 · similarity 0.86)