[Paper Review] A Tutorial on Parametric Variational Inference

Jens Sjölund|arXiv (Cornell University)|Jan 3, 2023

Genetic and phenotypic traits in livestock4 citations

TL;DR

This tutorial presents parametric variational inference as a scalable alternative to MCMC for Bayesian inference, focusing on optimizing a tractable lower bound (ELBO) to approximate intractable posteriors. It covers key techniques like the reparameterization trick, amortized inference via neural networks, and black-box variational inference with score function estimators, enabling efficient large-scale Bayesian modeling in modern machine learning.

ABSTRACT

Variational inference uses optimization, rather than integration, to approximate the marginal likelihood, and thereby the posterior, in a Bayesian model. Thanks to advances in computational scalability made in the last decade, variational inference is now the preferred choice for many high-dimensional models and large datasets. This tutorial introduces variational inference from the parametric perspective that dominates these recent developments, in contrast to the mean-field perspective commonly found in other introductory texts.

Motivation & Objective

• To provide a systematic introduction to parametric variational inference as a scalable alternative to MCMC for Bayesian inference.
• To explain how the evidence lower bound (ELBO) enables optimization-based approximation of intractable posterior distributions.
• To present modern techniques such as amortized variational inference and black-box variational inference for handling high-dimensional and large-scale models.
• To clarify the role of the Kullback-Leibler divergence in linking ELBO maximization to posterior approximation.
• To bridge the gap between theoretical foundations and practical implementation in modern probabilistic machine learning.

Proposed method

• Uses the evidence lower bound (ELBO) as a tractable surrogate for the intractable marginal likelihood, derived via Jensen's inequality and KL divergence decomposition.
• Employs the Kullback-Leibler divergence to measure the discrepancy between the true posterior and the variational approximation, with minimization equivalent to ELBO maximization.
• Applies the reparameterization trick to enable low-variance gradient estimation for differentiable variational families, facilitating stochastic gradient optimization.
• Introduces amortized variational inference via a neural network (encoder) that maps observations to variational parameters, reducing per-data-point optimization to a single model inference.
• Utilizes the score function estimator (REINFORCE) in black-box variational inference for non-reparameterizable models, with variance reduction techniques like control variates.
• Leverages automatic differentiation to compute gradients of the ELBO with respect to variational parameters, enabling end-to-end training in complex models.

Experimental results

Research questions

• RQ1 How can we approximate the intractable posterior distribution in Bayesian models using optimization instead of integration?
• RQ2 What is the role of the ELBO in variational inference, and how does it relate to the true marginal likelihood and KL divergence?
• RQ3 In what scenarios is the reparameterization trick applicable, and how does it improve gradient estimation in variational inference?
• RQ4 How does amortized variational inference reduce computational cost in large-scale Bayesian models with local latent variables?
• RQ5 What are the limitations of black-box variational inference, and how do variance reduction techniques improve its practicality?

Key findings

• The ELBO provides a lower bound on the log marginal likelihood that is both tractable and amenable to optimization, enabling scalable posterior approximation.
• Maximizing the ELBO is equivalent to minimizing the KL divergence between the variational posterior and the true posterior, ensuring the approximation becomes tighter.
• The reparameterization trick enables exact gradient computation via automatic differentiation, leading to low-variance gradient estimates in differentiable models.
• Amortized variational inference reduces per-data-point optimization to a single forward pass through a neural network, significantly improving scalability to large datasets.
• Black-box variational inference with score function estimators enables gradient-based optimization in non-reparameterizable models, albeit with higher variance that can be mitigated via control variates.
• The combination of amortization and reparameterization enables efficient training of modern models like variational autoencoders, where the encoder and decoder are jointly optimized via stochastic gradient descent.