[Paper Review] Trust-Region Variational Inference with Gaussian Mixture Models

Oleg Arenz, Mingjun Zhong+1 more|arXiv (Cornell University)|Jul 10, 2019

Gaussian Processes and Bayesian Inference70 references4 citations

TL;DR

This paper proposes Trust-Region Variational Inference with Gaussian Mixture Models (VIPS), a method that uses information-geometric trust regions and a lower-bound optimization to learn accurate, dynamically adapting GMM approximations of complex, multimodal target distributions. It achieves MCMC-quality samples with up to 1,000x less computational cost by enabling independent component updates and online component adaptation.

ABSTRACT

Many methods for machine learning rely on approximate inference from intractable probability distributions. Variational inference approximates such distributions by tractable models that can be subsequently used for approximate inference. Learning sufficiently accurate approximations requires a rich model family and careful exploration of the relevant modes of the target distribution. We propose a method for learning accurate GMM approximations of intractable probability distributions based on insights from policy search by using information-geometric trust regions for principled exploration. For efficient improvement of the GMM approximation, we derive a lower bound on the corresponding optimization objective enabling us to update the components independently. Our use of the lower bound ensures convergence to a stationary point of the original objective. The number of components is adapted online by adding new components in promising regions and by deleting components with negligible weight. We demonstrate on several domains that we can learn approximations of complex, multimodal distributions with a quality that is unmet by previous variational inference methods, and that the GMM approximation can be used for drawing samples that are on par with samples created by state-of-the-art MCMC samplers while requiring up to three orders of magnitude less computational resources.

Motivation & Objective

• Address the challenge of approximating complex, multimodal posterior distributions in Bayesian inference where standard variational methods fail due to restrictive mean-field assumptions.
• Overcome the limitations of MCMC and standard variational inference by combining computational efficiency with high sample quality for intractable posteriors.
• Develop a method that dynamically adapts the number of components in a Gaussian Mixture Model (GMM) during optimization to better capture multiple modes of the target distribution.
• Ensure stable and efficient optimization of the intractable KL divergence objective through a derived lower bound and trust-region updates.
• Enable high-quality sampling from the learned GMM approximation that matches or exceeds the performance of state-of-the-art MCMC samplers in sample quality while drastically reducing computational cost.

Proposed method

• Propose a trust-region framework inspired by policy search (MORE) to ensure stable, monotonic improvement of the variational objective via local quadratic approximations.
• Derive a lower bound on the reverse KL divergence that decomposes into independent sub-problems for each GMM component, enabling efficient, parallel updates.
• Use the lower bound to iteratively tighten and maximize the objective, guaranteeing convergence to a stationary point of the original KL minimization problem.
• Dynamically adjust the number of GMM components during optimization: add new components in high-density regions and remove components with negligible weights.
• Leverage function evaluations of the unnormalized target density to guide exploration, balancing exploitation of high-density regions and discovery of new modes.
• Apply a trust-region constraint to ensure that updates to the GMM parameters remain within a region where the local approximation is valid, preventing divergence.

Experimental results

Research questions

• RQ1 Can a trust-region optimization framework improve the stability and convergence of variational inference for multimodal posteriors using GMMs?
• RQ2 Can a lower-bound decomposition enable independent and efficient updates of GMM components while preserving convergence guarantees?
• RQ3 Can online component adaptation (addition and deletion) lead to more accurate and compact GMM approximations without prior knowledge of the number of modes?
• RQ4 How does the quality of samples from the learned GMM compare to those from state-of-the-art MCMC samplers in terms of statistical fidelity and computational cost?
• RQ5 To what extent can this method outperform existing variational inference baselines like BBVI and IAF in both ELBO and sample quality metrics on complex, multimodal distributions?

Key findings

• VIPS++ achieves maximum mean discrepancy (MMD) values comparable to state-of-the-art MCMC samplers like PTMCMC, with up to 1,000x less computational cost.
• The method learns GMM approximations that cover multiple modes of complex, high-dimensional posteriors (e.g., 20D GMM, planar robot with 4 goals) more effectively than BBVI and IAF.
• Despite learning a GMM with multiple components, IAF achieved a similar ELBO to VIPS++ on the planar robot task, but produced samples from only one of two main configurations, indicating poor mode coverage.
• The use of a trust-region framework ensures stable optimization and convergence to a stationary point of the original objective, even in high-dimensional, multimodal settings.
• Ground-truth samples for evaluation were generated using generalized elliptical slice sampling with up to two days of computation on 128 CPU cores, confirming the high quality of the baseline.
• On the German credit and breast cancer datasets, VIPS++ achieved significantly lower MMD values than BBVI and IAF, demonstrating superior sample quality across diverse probabilistic models.