PaperlayerView exact revision on arXiv

Source-linked AI summary

Automatic Differentiation Variational Inference

Alp Kucukelbir, Dustin Tran, Rajesh Ranganath, Andrew Gelman, David M. Blei

arXiv:1603.00788v1stat.MLcs.AIcs.LGstat.CO

TL;DR

Complex probabilistic models and large datasets make iterative modeling difficult because deriving and implementing scalable inference algorithms is laborious. ADVI automates variational inference by transforming latent variables into a common space and using Monte Carlo gradients, reparameterization, and stochastic optimization. The method was studied across ten models, applied to 1.7 million taxi trajectories, and deployed in Stan.

  • Problem

    Deriving and implementing model-specific variational inference algorithms is laborious, limiting iterative probabilistic modeling for complex models and large datasets.

  • Method

    ADVI transforms latent variables into a common space, computes ELBO gradients with Monte Carlo methods, reparameterizes them using a standard Gaussian, and optimizes the variational distribution.

  • Results

    ADVI was studied across ten probabilistic models and applied to 1.7 million taxi trajectories, with the model fit in minutes.

  • Takeaways & Limitations

    ADVI provides automated and scalable variational inference for a large class of probabilistic models and is available through Stan.

  • Takeaways & Limitations

    The current method uses first-order automatic differentiation and an adaptively tuned step-size heuristic; higher-order gradients or other optimization methods may improve convergence.

Abstract

from arXiv · show

Probabilistic modeling is iterative. A scientist posits a simple model, fits it to her data, refines it according to her analysis, and repeats. However, fitting complex models to large data is a bottleneck in this process. Deriving algorithms for new models can be both mathematically and computationally challenging, which makes it difficult to efficiently cycle through the steps. To this end, we develop automatic differentiation variational inference (ADVI). Using our method, the scientist only provides a probabilistic model and a dataset, nothing else. ADVI automatically derives an efficient variational inference algorithm, freeing the scientist to refine and explore many models. ADVI supports a broad class of models-no conjugacy assumptions are required. We study ADVI across ten different models and apply it to a dataset with millions of observations. ADVI is integrated into Stan, a probabilistic programming system; it is available for immediate use.

1 Introduction

ADVI automates variational inference for complex probabilistic models, addressing the computational bottleneck that limits iterative modeling with large and complex datasets. Implemented in Stan, it transforms models into a common space and uses automated differentiation and stochastic optimization to produce scalable inference algorithms.

  • Motivation: Probabilistic modeling supports an iterative cycle of proposing, fitting, analyzing, and refining models, but computation can prevent scientists from fully using it.Rich models for large, complex data require algorithms that are difficult to derive, implement, and scale.
  • Contribution: ADVI takes a probabilistic model and dataset as inputs and automatically produces a corresponding variational inference algorithm.It automates model-specific derivations and implementation rather than requiring the user to develop a custom optimization routine.
  • Applications: ADVI was applied across ten probabilistic modeling examples and to 1.7 million taxi trajectories, where the model could be fit in minutes.The taxi analysis used a non-conjugate mixture model with an unknown number of components.
  • Method: ADVI first transforms constrained latent variables into unconstrained real-valued variables, placing models in a common space for variational inference.The transformation incorporates the mapping into the joint distribution and removes the original latent-variable constraints.
  • Method: ADVI estimates variational gradients with Monte Carlo methods, reparameterizes them using a standard Gaussian, and optimizes with noisy gradients.An adaptively tuned step-size sequence supports convergence in practice.
  • Implementation: ADVI is implemented in Stan, which supplies transformations and automatic differentiation for computing gradients of the model’s log joint.Stan compiles a user-specified model into an inference executable.

2 Automatic Differentiation Variational Inference

ADVI automates variational inference by transforming differentiable probabilistic models into a common unconstrained space and optimizing a standardized variational approximation. Automatic differentiation and Monte Carlo integration produce stochastic gradients, extending automated inference to nonconjugate models.

  • 2.3 Automatic Transformation of Constrained Variables: ADVI transforms constrained latent variables into unconstrained real coordinates, allowing one variational family across models.For a positive latent rate, the logarithm maps R>0 to R and includes a Jacobian adjustment.
  • 2.1 Differentiable Probability Models: ADVI supports differentiable probability models with continuous latent variables and requires no conjugacy assumptions.The required gradient of the log joint must be valid within the prior support.
  • 2.2 Variational Inference: Variational inference minimizes KL divergence to the posterior by maximizing the ELBO, whose terms are expected log joint density and variational entropy.The ELBO equals negative KL divergence up to the constant log p(x).
  • 2.4 Variational Approximations in Real Coordinate Space: The transformed problem is unconstrained in the variational parameters, so ADVI can use gradient ascent without support-matching concerns.The parameter vector lives in an appropriately dimensioned real coordinate space.
  • 2.6 Stochastic Optimization: Automatic differentiation computes ELBO gradients, while Monte Carlo integration estimates expectations using standard-Gaussian samples.The resulting gradients are noisy and unbiased for any differentiable probability model; in practice, a single sample can suffice.
  • Related Work: Compared with earlier automated variational methods, ADVI supports a more comprehensive class of nonconjugate models.Earlier approaches described here relied on conjugacy, custom approximations, or expensive sampling.

3 Properties of Automatic Differentiation Variational Inference

ADVI’s accuracy depends on the variational approximation: mean-field estimates posterior means well but misses correlations and underestimates variances, while full-rank is more accurate at higher computational cost. Its stochastic gradients are efficient, but transformation choice also affects approximation quality.

  • Accuracy: Both ADVI variants correctly identify the analytic posterior mean for a correlated Gaussian, but mean-field produces an incorrect shape by ignoring covariance off-diagonal terms.The KL objective leads mean-field ADVI to systematically underestimate marginal variances.
  • Accuracy: MCMC and ADVI estimate logistic-regression posterior means similarly, while mean-field underestimates marginal posterior variances for most coefficients.The comparison uses 1000 simulated datapoints and Stan’s NUTS sampler.
  • Accuracy: Full-rank ADVI matches sampling for posterior mean estimates when stochastic volatility moves far from its mean, where mean-field struggles.The associated covariance analysis links these regions to stronger local posterior correlations.
  • Accuracy: Mean-field covariance fails to capture local correlation structure, whereas full-rank ADVI and sampling recover correlations where log volatility trends away from its mean.All covariance estimates show blurry spread because of finite sample size.
  • Recommendations: Scientists interested in posterior variances and covariances should use full-rank ADVI, but its computational cost can be prohibitive for large datasets.Full-rank captures posterior correlations and thereby produces more accurate marginal variance estimates.
  • Stochastic gradients: ADVI gradient estimates are unbiased, and experiments find ADVI gradients statistically more efficient than BBVI gradients across univariate and 100-dimensional nonlinear-regression models.In the univariate example, a single Monte Carlo sample suffices in practice because the ADVI gradient has lower variance.
  • Sensitivity to transformations: Both graphical and numerical results prefer transformation T2 over T1 because T2 is nearly linear for large θ while the logarithm T1 flattens.The posterior and ADVI approximation are both light-tailed in this comparison.

4 Automatic Differentiation Variational Inference in Practice

ADVI applies automatically derived variational inference to diverse nonconjugate models, enabling rapid comparisons with MCMC and analyses of large datasets. Across examples, it achieves competitive or faster inference and supports iterative exploration of millions of taxi trajectories.

  • Models evaluated: ADVI was evaluated on linear regression, hierarchical logistic regression, matrix factorization, mixture models, and probabilistic principal component analysis.The study used simulated and real data across an array of nonconjugate probability models.
  • Regression models: For the regression models, all methods reached the same predictive accuracy, while one Monte Carlo sample per iteration was sufficient and fastest for ADVI.The result used average log predictive accuracy over time and set M = 1 thereafter.
  • Matrix factorization: ADVI provided an order of magnitude speed improvement over NUTS for both non-negative matrix factorization models, while HMC produced no useful samples within one hour.NUTS struggled particularly with the Dirichlet Exponential model.
  • Gaussian mixture models: For the imageCLEF mixture model, ADVI quickly found a good solution while NUTS struggled and HMC failed, likely because label switching affects HMC-based mixture inference.The experiment used 1,000 images and 10 mixture components.
  • Gaussian mixture models: On the full imageCLEF dataset, ADVI reached high predictive accuracy with minibatches of 500 or larger and converged in about two hours, whereas NUTS could not handle the dataset.Smaller minibatches led to suboptimal solutions.
  • Taxi trajectory case study: For 1.7 million taxi trajectories, ADVI found an eleven-dimensional subspace in approximately an hour and produced geographical clusters of frequently taken routes.The trajectory analysis used a 10,000-trajectory subsample for subspace inference and 50,000 randomly sampled trajectories for visualization.
  • Taxi trajectory case study: Interpolating every taxi trajectory to equal length discarded duration information, limiting the case study’s ability to assess whether roads lead to longer trips.The authors explicitly raise this as a question for further analysis.
  • Taxi trajectory case study: ADVI enabled a revised taxi analysis to find more informative Gaussian mixture clusters in less than two hours, identifying busy bridges in Porto.The clusters distinguished short routes using old city-center bridges from longer routes using newer highway connections.

5 Discussion

The discussion presents ADVI as a broadly applicable variational inference tool deployed in Stan, while identifying accuracy, optimization, heuristic, and model-class limitations for future work.

  • Contribution: ADVI transforms latent variables into a common space, supports a large class of probabilistic models, and was studied across ten models before deployment in Stan.The common-space formulation allows solving the variational inference problem across the model class.
  • Accuracy: ADVI’s accuracy can be sensitive to the transformations mapping constrained parameters into real coordinates.The discussion points to richer transformation families as possible ways to improve accuracy.
  • Optimization: ADVI uses first-order automatic differentiation for stochastic gradient ascent, while higher-order gradients, line search, and natural gradients may improve convergence at added computational cost or complexity.The discussion presents these as research directions rather than established improvements.
  • Practical heuristics: Initialization and step-size scaling affect convergence, and current step-size tuning uses a finite search that adds computation.The authors suggest model- and dataset-adaptive heuristics based on moment matching.
  • Scope: Stan-focused ADVI targets differentiable probability models, leaving discrete latent variables as an extension that may require higher-variance gradient estimators.The proposed extension is described as requiring care.

A Transformations of Continuous Probability Densities

A differentiable one-to-one transformation changes a probability density through the derivative of its inverse, generalized multivariately by the absolute Jacobian determinant.

  • Univariate transformation: For a one-to-one differentiable transformation, the transformed density is obtained by applying the inverse transformation and its absolute derivative.The proof follows from transforming cumulative densities and differentiating them.
  • Multivariate transformation: In multiple dimensions, the corresponding density adjustment is the absolute determinant of the inverse transformation’s Jacobian.The Jacobian generalizes the absolute derivative from scalar to multivariate variables.
  • Interpretation: The Jacobian accounts for how transformations warp unit volumes so that the transformed density still integrates to one.For linear and affine transformations that preserve unit volumes, the Jacobian adjustment evaluates to one.

B Transformation of the Evidence Lower Bound

ADVI begins with the evidence lower bound in the original latent-variable space and transforms the latent variables into real coordinate space.

  • Transformed variational family: The variational approximation is represented in real coordinate space as q(ζ; φ), with ζ = T(θ).The transformation links the original latent variables θ to the real-coordinate variables ζ.
  • ELBO transformation: ADVI first formulates the ELBO in the original latent-variable space and then transforms that space into real coordinates.This transformation is the stated sequence for changing the variational inference problem’s coordinate system.

C Gradients of the Evidence Lower Bound

The evidence lower bound gradients are derived separately for the µ, mean-field ω, and full-rank L parameters. The µ derivation exchanges differentiation and integration using dominated convergence, then applies the chain rule.

  • The gradient with respect to µ exchanges differentiation and integration through the dominated convergence theorem, followed by the chain rule.This establishes the gradient calculation for the µ parameter.
  • The mean-field ω parameter has its gradient considered separately.
  • The full-rank L parameter has its gradient considered separately.

D Automating Expectations: Monte Carlo Integration

ADVI approximates the expectation terms in the evidence lower bound with Monte Carlo integration. This requires samples from q and produces noisy but unbiased integral estimates.

  • Expectations are approximated with Monte Carlo integration using samples from q.
  • Monte Carlo integration produces estimates that are noisy yet unbiased.
  • The standard deviation of the estimates is described as being of order 1/.

E Running ADVI in Stan

ADVI is available through Stan, including its command-line interface, cmdStan. The workflow uses a dataset file, produces posterior samples, and reports the ELBO.

  • Stan provides multiple interfaces, including the command-line interface cmdStan.
  • The cmdStan usage syntax for ADVI is presented in Figure 16.
  • Data.data.R stores the dataset in R dump format, while output_advi.csv contains posterior samples and elbo_advi.csv reports the ELBO.

F.1 Linear Regression with Automatic Relevance Determination

The paper demonstrates ADVI on linear regression with automatic relevance determination, a high-dimensional sparse regression model. The model uses Gaussian-noise observations and independent Gamma priors for weight-specific hyperparameters, with simulated data containing predictive and non-predictive regressions.

  • Linear regression with ARD is presented as a high-dimensional sparse regression model.The corresponding Stan code is provided in Figure 17.
  • The inputs x_n are D-dimensional, the outputs y_n are one-dimensional, and the weights vector w is D-dimensional.
  • The likelihood models measurements corrupted by iid Gaussian noise with unknown standard deviation σ.
  • Each component of the D-dimensional hyper-prior α receives its own independent Gamma prior.
  • The simulated data are constructed so that only half of the regressions have predictive power.Figure 10 uses a0 = b0 = c0 = d0 = 1 as the Gamma-prior hyper-parameters.

F.2 Hierarchical Logistic Regression

The section presents several probabilistic models implemented in Stan, including hierarchical logistic regression, matrix factorization, Gaussian mixtures, and principal-component models. It also documents model-specific priors, identifiability constraints, sparsity choices, and ADVI implementations.

  • Hierarchical Logistic Regression: Hierarchical logistic regression models voting preferences using structured predictors and hierarchical variables, with standard deviation hyper-priors constrained between 0 and 100.The model is motivated by a dataset from the 1988 United States presidential election.
  • Gamma Poisson Factorization: The Gamma Poisson factorization model represents discrete user–item matrices with K factors, but swapping factor rows and columns leaves the inner product unchanged.The experiments constrain each user factor vector to be ordered, set K = 10, and use Gamma hyper-parameters equal to 1.
  • Dirichlet Exponential Factorization: The Dirichlet Exponential model combines Dirichlet-enforced uniqueness with exponential-induced sparsity in a non-conjugate model.With K = 10, α0 = 1000, and λ0 = 0.1, the factor vectors β_i appear sparse.
  • Gaussian Mixture Model: The Gaussian mixture model groups natural images by color histograms using Gaussian, lognormal, and Dirichlet priors for means, standard deviations, and mixture components.The imageCLEF color histograms have D = 576 dimensions, and α0 = 1000 is used in the experiments.
  • Probabilistic Principal Component Models: PPCA introduces latent M-dimensional variables and principal components with standard-normal priors, while ARD adds a vector α that selects retained components.SUP-PPCA additionally regresses observed variables y onto the principal-component subspace to support prediction.
Loading 1603.00788v1…