[Paper Review] Automated Variational Inference in Probabilistic Programming

David Wingate, Théophane Weber|arXiv (Cornell University)|Jan 7, 2013

Bayesian Modeling and Causal InferenceComputer Science26 references72 citations

TL;DR

This paper introduces an automated variational inference algorithm for probabilistic programming that uses stochastic gradient optimization to efficiently approximate posterior distributions without requiring analytical tractability. By leveraging gradient estimation through trace-based computation and reinforcement learning-inspired baseline adaptation, the method achieves faster convergence and lower variance than standard stochastic gradient descent, especially in complex, non-conjugate models like LDA and QMR-DT.

ABSTRACT

We present a new algorithm for approximate inference in probabilistic programs, based on a stochastic gradient for variational programs. This method is efficient without restrictions on the probabilistic program; it is particularly practical for distributions which are not analytically tractable, including highly structured distributions that arise in probabilistic programs. We show how to automatically derive mean-field probabilistic programs and optimize them, and demonstrate that our perspective improves inference efficiency over other algorithms.

Motivation & Objective

• Address the challenge of performing efficient, automated variational inference in arbitrary probabilistic programs where analytical solutions are intractable.
• Overcome the limitations of traditional coordinate ascent variational inference, which requires manual derivation of update equations and fails for non-conjugate or highly structured models.
• Develop a general-purpose, automatic method that can derive and optimize mean-field variational approximations directly from probabilistic programs without requiring model-specific derivations.
• Improve inference efficiency and convergence speed by using stochastic gradient optimization with baseline adaptation, inspired by reinforcement learning.
• Enable scalable posterior approximation in complex models such as LDA and QMR-DT, where standard methods struggle due to high dimensionality and non-conjugacy.

Proposed method

• Formalize probabilistic programs as trace-based generative processes, where each execution trace corresponds to a sequence of elementary random procedures (ERPs) with history-dependent distributions.
• Define the variational objective as maximizing the evidence lower bound (ELBO), which is equivalent to minimizing the KL divergence between the variational approximation and the true posterior.
• Derive a stochastic gradient estimator for the ELBO using pathwise derivatives and reparameterization, enabling gradient computation through stochastic nodes in the program.
• Introduce a baseline adaptation mechanism (inspired by the Episodic Natural Actor-Critic) to reduce gradient variance, improving optimization stability and convergence speed.
• Use online, mini-batch estimation of the gradient to allow scalable optimization over large datasets without reprocessing the entire data set.
• Apply the method to both LDA and QMR-DT benchmarks, demonstrating its ability to automatically generate and optimize variational programs without manual derivation.

Experimental results

Research questions

• RQ1 Can stochastic gradient-based variational inference be automatically applied to arbitrary probabilistic programs without requiring analytical derivations of update equations?
• RQ2 How does the use of a learned baseline in gradient estimation affect convergence speed and variance in variational inference for complex models?
• RQ3 Can the proposed method outperform standard stochastic gradient descent and second-order methods in terms of convergence and stability on non-conjugate models like LDA and QMR-DT?
• RQ4 To what extent can automated variational inference scale to large datasets and complex, structured models where traditional methods fail?
• RQ5 Does the automatic derivation of mean-field variational programs lead to accurate and efficient posterior approximations suitable for downstream inference tasks?

Key findings

• The proposed automated variational inference (AVI) method achieves faster convergence and lower gradient variance compared to standard stochastic gradient descent on both LDA and QMR-DT benchmarks.
• The Episodic Natural Actor-Critic (ENAC) algorithm, which includes baseline adaptation, significantly outperforms vanilla stochastic gradient descent, even when both use the same step size and gradient normalization.
• Second-order gradient descent (SOGD), despite using curvature information, performs poorly and even diverges in the LDA experiment, indicating that curvature estimation alone is insufficient without variance reduction.
• Conjugate gradient optimization with ENAC gradients converges faster than steepest descent, demonstrating that the quality of the gradient direction is the key factor, not just second-order information.
• Once trained, the variational program enables fast posterior sampling at a fraction of the cost of MCMC, and samples can be used as warm starts for MCMC to further improve efficiency.
• The method successfully automates the derivation and optimization of mean-field variational approximations in complex, non-conjugate models, eliminating the need for manual derivation of coordinate ascent equations.