[Paper Review] Automatic Variational Inference in Stan

Alp Kucukelbir, Rajesh Ranganath+2 more|arXiv (Cornell University)|Jun 10, 2015

Gaussian Processes and Bayesian InferenceComputer Science12 references71 citations

TL;DR

This paper introduces Automatic Differentiation Variational Inference (ADVI), an automated method for variational inference in Stan that requires only a probabilistic model and data. By transforming latent variables to the real space, using a mean-field Gaussian variational family, and leveraging automatic differentiation and stochastic optimization, ADVI enables scalable Bayesian inference across diverse models, achieving orders-of-magnitude speedups over MCMC on large datasets like 250,000 images.

ABSTRACT

Variational inference is a scalable technique for approximate Bayesian inference. Deriving variational inference algorithms requires tedious model-specific calculations; this makes it difficult to automate. We propose an automatic variational inference algorithm, automatic differentiation variational inference (ADVI). The user only provides a Bayesian model and a dataset; nothing else. We make no conjugacy assumptions and support a broad class of models. The algorithm automatically determines an appropriate variational family and optimizes the variational objective. We implement ADVI in Stan (code available now), a probabilistic programming framework. We compare ADVI to MCMC sampling across hierarchical generalized linear models, nonconjugate matrix factorization, and a mixture model. We train the mixture model on a quarter million images. With ADVI we can use variational inference on any model we write in Stan.

Motivation & Objective

• To automate variational inference in probabilistic programming frameworks, eliminating the need for model-specific derivation of variational families and objectives.
• To support a broad class of differentiable probability models, including non-conjugate and hierarchical models, without requiring conjugacy assumptions.
• To enable scalable inference on large datasets by integrating stochastic variational inference with automatic differentiation.
• To provide a plug-and-play inference solution within Stan, allowing users to apply variational inference to any model written in Stan’s language.
• To achieve performance comparable to MCMC but orders of magnitude faster, especially on massive datasets.

Proposed method

• Transform latent variables to the real space using monotonic transformations (e.g., log for positive variables), enabling the use of a standard multivariate normal variational family.
• Use a mean-field Gaussian approximation in the transformed space, which induces a non-Gaussian approximation in the original parameter space.
• Leverage automatic differentiation to compute gradients of the variational lower bound (ELBO) with respect to variational parameters, enabling efficient optimization.
• Apply stochastic optimization with mini-batches to scale to large datasets, using a scaling factor to correct for subsampling bias in the ELBO.
• Integrate the method into Stan’s probabilistic programming framework, allowing users to write models in Stan’s language and automatically apply ADVI without additional implementation.
• Use reparameterization gradients and black-box variational inference techniques to optimize the ELBO efficiently without requiring analytical derivations.

Experimental results

Research questions

• RQ1 Can variational inference be fully automated in a general-purpose probabilistic programming framework like Stan without requiring model-specific derivations?
• RQ2 Can automatic differentiation and stochastic optimization enable scalable Bayesian inference on large datasets, such as 250,000 images, where MCMC is infeasible?
• RQ3 Does the proposed automatic variational inference method achieve competitive accuracy compared to MCMC while being significantly faster?
• RQ4 Can the method handle a broad class of models, including non-conjugate and hierarchical models, without requiring conjugacy assumptions?
• RQ5 How does the performance of ADVI compare to Stan’s default MCMC sampler (NUTS) in terms of speed and predictive accuracy?

Key findings

• ADVI achieves orders-of-magnitude speedup over NUTS, Stan’s default MCMC sampler, on a hierarchical logistic regression model with 1,000 images.
• On a non-conjugate Gaussian mixture model trained on 250,000 images, ADVI produced a valid posterior approximation in under two hours, a task that is infeasible with standard MCMC.
• ADVI outperformed NUTS in held-out predictive likelihood on the 1,000-image dataset, demonstrating competitive model fitness despite being faster.
• The method successfully supports a wide range of models, including hierarchical generalized linear models, non-conjugate matrix factorization, and Gaussian mixture models, without requiring model-specific modifications.
• Stochastic variational inference with subsampling enabled ADVI to scale to large datasets, maintaining accuracy while reducing computation time.
• ADVI was successfully integrated into Stan 2.7 and later, providing a seamless, automated inference option for any Stan model.