Variational Inference with Mixtures of Isotropic Gaussians

Abstract: Variational inference (VI) is a popular approach in Bayesian inference, that looks for the best approximation of the posterior distribution within a parametric family, minimizing a loss that is typically the (reverse) Kullback-Leibler (KL) divergence. In this paper, we focus on the following parametric family: mixtures of isotropic Gaussians (i.e., with diagonal covariance matrices proportional to the identity) and uniform weights. We develop a variational framework and provide efficient algorithms suited for this family. In contrast with mixtures of Gaussian with generic covariance matrices, this choice presents a balance between accurate approximations of multimodal Bayesian posteriors, while being memory and computationally efficient. Our algorithms implement gradient descent on the location of the mixture components (the modes of the Gaussians), and either (an entropic) Mirror or Bures descent on their variance parameters. We illustrate the performance of our algorithms on numerical experiments.

• The paper presents a new variational framework using mixtures of isotropic Gaussians to balance accuracy and computational efficiency in approximating Bayesian posteriors.
• It introduces two novel optimization schemes—IBW and MD—for updating component means and variances, offering linear computational savings over full-covariance models.
• Empirical evaluations on synthetic and real-world datasets demonstrate comparable accuracy to full models while effectively mitigating mode collapse.

Variational Inference with Mixtures of Isotropic Gaussians

This paper presents a study on variational inference (VI) applying mixtures of isotropic Gaussians as the variational family. The work addresses the balance between accuracy and computational efficiency in approximating Bayesian posterior distributions, especially for scenarios involving multimodal distributions. The authors develop a variational framework that leverages isotropic covariance matrices, enhancing memory and computational feasibility while maintaining representation capabilities for complex distributions.

Methodology

The authors propose a novel approach within the VI framework by focusing on mixtures of isotropic Gaussians, characterized by diagonal covariance matrices proportional to the identity matrix, and employing uniform weights across components. This choice of variational family presents a promising trade-off: while it is less flexible than full covariance Gaussian mixtures, it is computationally efficient and sufficiently expressive to approximate multimodal distributions.

Two primary optimization schemes for minimizing the reverse Kullback-Leibler (KL) divergence are introduced:

• Gradient Descent: Applied to optimize the locations of the mixture components, performing standard gradient descent on the component means.
• Variance Optimization: Two approaches are explored—(i) Bures-Wasserstein (IBW) update employing a Riemannian gradient flow dictated by the Bures metric (appropriate for covariance matrices) and (ii) Entropic Mirror Descent (MD) update enforcing positive variance updates efficiently.

Key Contributions

1. Variational Framework and Algorithms: The development of algorithms specifically tailored to mixtures of isotropic Gaussians allows practitioners to efficiently approximate complex posterior distributions, particularly those displaying multimodality.
2. Empirical Evaluation: The authors provide a comprehensive empirical evaluation across both synthetic and real-world datasets, demonstrating the practical applicability and benefits of the proposed approach. The experiments are designed to assess both accuracy in terms of KL divergence and computational efficiency.
3. Balance between Computational Efficiency and Accuracy: The work employs a detailed analysis of the trade-offs involved in choosing isotropic Gaussians over more traditional full-covariance Gaussian mixtures—striking an appropriate balance between modeling power and computational resource demands.
4. Algorithmic Innovations: The introduction of IBW and MD schemes showcases new ways to navigate the trade-off between computational cost (linear compared to quadratic costs in full covariance models) and approximation quality.

Numerical and Experimental Insights

The experimental results reveal that the proposed methods (IBW and MD) often achieve similar accuracy to full-covariance models but with significantly reduced computational load, especially in higher dimensions. Additionally, issues such as mode collapse, which can emerge in other VI approaches like normalizing flows, are effectively mitigated through the isotropic Gaussian mixture strategy. Furthermore, the inclusion of benchmark comparisons with existing methods such as Normalizing Flows, Hamiltonian Monte Carlo, and Automatic Differentiation VI situates the work within the broader landscape of approximate inference techniques.

Implications and Future Work

The practical implications of this research span across various machine learning and statistics applications where Bayesian inference is crucial but traditional MCMC or full variational methods face limitations due to scalability or execution speed. The efficiency of isotropic Gaussian mixtures paves the way for handling larger datasets and more complex models in a feasible manner.

Future directions may involve extending this framework to incorporate more sophisticated models, improving convergence analysis, or exploring adaptive methods to dynamically tune the number of Gaussian components based on the target posterior's complexity. Further theoretical analysis could also be beneficial in quantifying the gap in approximation accuracy between isotropic and full-covariance mixture models.

In summary, this study advances the variational inference domain by crafting a structured movement towards computationally efficient approximate inference, with robust theoretical underpinning and validated across diverse empirical scenarios.