Source-linked AI summary
Black Box Variational Inference
Rajesh Ranganath, Sean Gerrish, David M. Blei
TL;DR
Approximate posterior inference often requires tedious model-specific derivations, limiting rapid exploration of complex models. The paper introduces Black Box Variational Inference, which optimizes a variational objective with Monte Carlo stochastic gradients and model-independent variance reduction. It reaches better predictive likelihoods much faster than sampling methods and supports rapid exploration of longitudinal healthcare models.
Problem
Generic variational inference often lacks closed-form solutions and requires model-specific computations, making it difficult to rapidly develop and revise models.
Method
Black Box Variational Inference uses stochastic optimization with Monte Carlo gradients from the variational distribution, combined with variance-reduction methods that avoid model-specific derivations.
Results
The method reaches better predictive likelihoods much faster than sampling methods and quickly fits several longitudinal patient-data models.
Takeaways & Limitations
Black Box Variational Inference reduces analytic burden and works well for new models while requiring minimal analytic work from practitioners.
Takeaways & Limitations
Future improvements include broader score-function library support, dynamic sample counts, and sampling methods that could reduce variance.
Abstract
from arXiv · showhide
Variational inference has become a widely used method to approximate posteriors in complex latent variables models. However, deriving a variational inference algorithm generally requires significant model-specific analysis, and these efforts can hinder and deter us from quickly developing and exploring a variety of models for a problem at hand. In this paper, we present a "black box" variational inference algorithm, one that can be quickly applied to many models with little additional derivation. Our method is based on a stochastic optimization of the variational objective where the noisy gradient is computed from Monte Carlo samples from the variational distribution. We develop a number of methods to reduce the variance of the gradient, always maintaining the criterion that we want to avoid difficult model-based derivations. We evaluate our method against the corresponding black box sampling based methods. We find that our method reaches better predictive likelihoods much faster than sampling methods. Finally, we demonstrate that Black Box Variational Inference lets us easily explore a wide space of models by quickly constructing and evaluating several models of longitudinal healthcare data.
1 Introduction
Black Box Variational Inference targets the model-specific derivations that make approximate posterior inference difficult to apply and slow model exploration. It uses Monte Carlo stochastic gradients, variance reduction, and generic optimization to support rapid application across models.
- Motivation: Generic variational inference lacks closed-form solutions for arbitrary models and families, often requiring model-specific computations.For many models, the required expectations are intractable, making derivations tedious and limiting practical use.
- Approach: Black Box Variational Inference rewrites the variational objective gradient as an expectation estimated from samples of the variational distribution.The sampled gradients are used in stochastic optimization of the variational parameters.
- Practicality: Practitioners need only evaluate model log-likelihood-related functions, while sampling and Monte Carlo calculations can be shared across models.This reduces the additional implementation work needed for new modeling settings.
- Variance and scaling: Rao-Blackwellization and control variates reduce gradient-estimator variance without requiring model-specific computations.The paper also uses adaptive learning rates and generic stochastic variational inference with data subsampling.
- Evaluation: The method reaches better predictive likelihoods much faster than sampling methods.The authors compare it with Metropolis-Hastings-in-Gibbs and apply it to several longitudinal patient-data models.
2 Black Box Variational Inference
The method turns ELBO optimization into stochastic optimization using unbiased Monte Carlo gradient estimates that require only reusable variational-distribution operations and model log-joint evaluation.
- Objective: Variational inference approximates a conditional distribution by optimizing a simple variational family toward the closest distribution in KL divergence.Maximizing the ELBO is equivalent to minimizing this KL divergence.
- Motivation: Closed-form coordinate updates are limited to conditionally conjugate models and require analytic expectations for each new model.Outside these settings, the bookkeeping becomes tedious and model-specific.
- Algorithm: The algorithm estimates the ELBO gradient by sampling latent variables, evaluating the joint and variational distributions, and applying the variational score function.The resulting noisy but unbiased gradients update the variational parameters through stochastic optimization.
- Optimization: The procedure uses a Robbins-Monro learning-rate sequence to iteratively update variational parameters until the parameter change is below 0.01.Algorithm 1 draws S samples from the variational distribution at each iteration.
- Black-box property: Score functions and sampling procedures depend only on the variational distribution and can be derived once and reused across models.The model-specific requirement is computing the joint log probability for sampled variables.
3 Controlling the Variance
The paper controls the high variance of Monte Carlo ELBO gradients with Rao-Blackwellization and control variates, then combines them into a reusable black-box algorithm.
- Motivation: High-variance gradients can force very small optimization steps and slow convergence, motivating model-independent variance reduction.The paper develops Rao-Blackwellization and control variates without model-specific derivations.
- 3.1 Rao-Blackwellization: Rao-Blackwellization replaces a random variable with a conditional expectation, producing a lower-variance estimator while preserving its expectation.For gradient components, the construction avoids model-specific conditional-expectation computations.
- 3.1 Rao-Blackwellization: The Rao-Blackwellized gradient estimator computes each component using variables in the relevant Markov blanket and factorized variational structure.The approach applies to mean-field and other factored variational approximations.
- 3.2 Control Variates: Control variates preserve the target expectation while reducing Monte Carlo variance, with scaling chosen from covariance and variance estimates.The score function of the variational approximation is used because its expectation is zero and it depends only on the variational distribution.
- 3.2 Control Variates: The combined estimator uses Rao-Blackwellization and per-component control variates to further reduce gradient variance.Black Box Variational Inference then feeds these noisy gradients into stochastic optimization of the ELBO.
- 3.3 Black Box Variational Inference (II): Algorithm 2 requires samples from the variational approximation, variational-distribution computations, and easy model computations.The algorithm can store samples for simplicity, while the method is intended for use across many models.
4 Extensions
The paper extends black-box variational inference with adaptive per-component learning rates and observation subsampling for scalable stochastic gradients.
- AdaGrad: AdaGrad sets per-component learning rates using the diagonal of accumulated noisy-gradient outer products.The diagonal matches the gradient dimension and captures noise and varying length scales.
- AdaGrad: The adaptive schedule reduces the algorithm’s learning-rate parameters relative to the standard two-parameter Robbins–Monro schedule.
- Stochastic Inference in Hierarchical Bayesian Models: The hierarchical model contains hyperparameter η, global latent variables β, local latent variables z1...n, and observations x1...n.
- Stochastic Inference in Hierarchical Bayesian Models: Observation subsampling scales stochastic variational inference by replacing iteration over every datapoint with sampled observations and variational-distribution samples.This construction applies to hierarchical Bayesian models under the mean field approximating family.
5 Empirical Study
The empirical study evaluates Black Box Variational Inference on longitudinal healthcare data, comparing its convergence, variance-reduction strategies, and model-exploration capabilities with sampling-based methods.
- Data and evaluation: The study uses longitudinal chronic-kidney-disease data from 976 patients, with 803 training and 173 test patients and 33K visits.Each visit contains a subset of 17 laboratory measurements, collected at sparse and irregularly spaced times.
- Data and evaluation: The modeling goal is to obtain low-dimensional latent factors representing patient health from laboratory measurements at each visit.The factors are positive random variables intended to summarize visits.
- Models: The Gamma-Normal time-series model uses normal weights and observations, while visit-level factors are drawn from a gamma distribution.The model propagates visit factors through time and permits factors to affect laboratory measurements positively or negatively.
- Method comparison: Black Box Variational Inference outperformed Metropolis-Hastings within Gibbs on predictive likelihood and reached better values faster under a fixed 20-hour computational budget.The comparison used held-out test data and showed similar results across different random initializations; the Gibbs sampler’s progress slowed considerably after 5 hours.
- Variance reductions: Rao-Blackwellization reduced gradient-estimator variance by several orders of magnitude, and control variates reduced it further.The algorithm without variance reduction failed to make noticeable progress within the allotted time, whereas the reductions improved convergence speed.
- Model exploration: Black Box Variational Inference enabled rapid fitting and comparison of four nonconjugate models, while Gamma-weight models performed relatively poorly on the patient dataset.The paper attributes the poorer Gamma-model performance to their inability to capture negative correlations among laboratory measurements.
6 Conclusion
The paper presents Black Box Variational Inference as a stochastic ELBO optimization method that reduces analytic burden through sampling and model-free variance reduction. It works well for new models with minimal analytic effort, while several sampling and software improvements remain open.
- Black Box Variational Inference drastically reduces the analytic burden of variational inference.The algorithm is designed to require minimal analytic work from practitioners.
- The method optimizes the ELBO stochastically by sampling from the variational posterior to compute a noisy gradient.Model-free variance reductions are essential to reducing the variance of that gradient.
- Black Box Variational Inference works well for new models while requiring minimal analytic work by the practitioner.
- Future improvements include broader score-function library support, dynamically setting the number of samples, and quasi-Monte Carlo sampling.The authors expect carefully selected samples to significantly decrease sampling variance.
7 Appendix: The Gradient of the ELBO
The appendix derives the ELBO gradient as an expectation under the variational distribution. The derivation relies on exchanging derivatives and integrals under stated regularity conditions and rewriting the variational-density derivative with its score function.
- The ELBO gradient is derived by differentiating the objective and expressing its terms through expectations under the variational distribution.
- The derivation exchanges derivatives with integrals using the dominated convergence theorem.It also uses that ∇λ[log p(x, z)] = 0.
- The first term in the differentiated ELBO expression is zero.
- The identity ∇λ[q(z|λ)] = ∇λ[log q(z|λ)]q(z|λ) converts the remaining derivative into an expectation involving the variational score.
- The derivation assumes that the score function exists and that the score and likelihoods are bounded.
Supplement
The supplement develops Rao-Blackwellized and subsampled noisy-gradient estimators for factorized variational distributions. These estimators integrate out variables where possible, can use observation subsampling, and reduce variance with control variates.
- The Rao-Blackwellized estimator is presented for the full mean-field family and also applies to factorizations such as structured mean-field.
- Rao-Blackwellization lowers estimator variance by integrating out variables under a factorized joint distribution.For each gradient component, expectations are computed with respect to the other factors.
- Under the mean-field assumption, the gradient for a factor depends on expectations over variables in that factor’s complete conditional.Variables outside the latent variable’s Markov blanket can be Rao-Blackwellized without model-specific computations.
- The supplement situates these estimators in hierarchical Bayesian models with global and local variational parameters under a mean-field approximation.
- Observation subsampling produces an unbiased noisy ELBO gradient without iterating over every observation at each update.The estimator samples an observation uniformly and uses the expected-value property to define stochastic optimization.
- Control variates and optimal scalings are used to construct a lower-variance noisy gradient.