Gaussian Mixture Variational Inference with Wasserstein–Fisher–Rao Gradient Flows

The choice of geometry affects the efficiency of gradient-based algorithms for variational inference (VI). We study Wasserstein–Fisher–Rao (WFR) gradient flows over Gaussians and their mixtures. For Gaussian VI with smooth, strongly log-concave targets, we prove exponential convergence of a WFR splitting scheme to the optimal Gaussian approximation. For general targets, we establish stationarity guarantees for Gaussian and mixture VI. The splitting scheme evaluates target statistics twice per iteration, once for each geometry. For efficiency, we instead develop a switched algorithm that chooses one suitable geometry per update. For Gaussian targets, we show that Wasserstein warm-up followed by Fisher–Rao refinement avoids the logarithmic initialization delay of Fisher–Rao and the slow convergence of Wasserstein on ill-conditioned targets. We extend this approach to Gaussian mixtures and propose adaptive rules for choosing each component’s switching time. We prove descent and stationarity guarantees for the resulting algorithm. Across benchmarks in machine learning and PDE-based Bayesian inverse problems, the switched method is competitive with or better than competing methods and is substantially less sensitive to initialization.