Automatic Variational Inference in Stan

Introduction

Bayesian inference is a powerful framework for analyzing data. We design a model for data using latent variables; we then analyze data by calculating the posterior density of the latent variables. For machine learning models, calculating the posterior is often difficult; we resort to approximation.

Variational inference (vi) approximates the posterior with a simpler density . We search over a family of simple densities and find the member closest to the posterior. This turns approximate inference into optimization. vi has had a tremendous impact on machine learning; it is typically faster than Markov chain Monte Carlo (mcmc) sampling (as we show here too) and has recently scaled up to massive data .

Unfortunately, vi algorithms are difficult to derive. We must first define the family of approximating densities, and then calculate model-specific quantities relative to that family to solve the variational optimization problem. Both steps require expert knowledge. The resulting algorithm is tied to both the model and the chosen approximation.

In this paper we develop a method for automating variational inference, automatic differentiation variational inference (advi). Given any model from a wide class (specifically, differentiable probability models), advi determines an appropriate variational family and an algorithm for optimizing the corresponding variational objective. We implement advi in Stan , a flexible probabilistic programming framework originally designed for sampling-based inference. Stan describes a high-level language to define probabilistic models (e.g., Figure 2) as well as a model compiler, a library of transformations, and an efficient automatic differentiation toolbox. With advi we can now use variational inference on any model we can express in Stan.advi is available in Stan 2.7 (development branch). It will appear in Stan 2.8. See Appendix C. (See Appendices F to J.)

Figure 1 illustrates the advantages of our method. We present a nonconjugate Gaussian mixture model for analyzing natural images; this is 40 lines in Stan (Figure 10). Section 1(a) illustrates Bayesian inference on $1000$ images. The $y$-axis is held-out likelihood, a measure of model fitness; the $x$-axis is time (on a log scale). advi is orders of magnitude faster than nuts, a state-of-the-art mcmc algorithm (and Stan’s default inference technique) . We also study nonconjugate factorization models and hierarchical generalized linear models; we consistently observe speed-up against nuts.

Section 1(b) illustrates Bayesian inference on $250\,000$ images, the size of data we more commonly find in machine learning. Here we use advi with stochastic variational inference , giving an approximate posterior in under two hours. For data like these, mcmc techniques cannot even practically begin analysis, a motivating case for approximate inference.

Related Work. advi automates variational inference within the Stan probabilistic programming framework . This draws on two major themes.

The first is a body of work that aims to generalize vi. Kingma and Welling and Rezende et al. describe a reparameterization of the variational problem that simplifies optimization. Ranganath et al. and Salimans and Knowles propose a black-box technique that only uses the gradient of the approximating family for optimization. Titsias and Lázaro-Gredilla leverage the gradient of the model for a small class of models. We build on and extend these ideas to automate variational inference; we highlight technical connections as we develop our method.

The second theme is probabilistic programming. Wingate and Weber study vi in general probabilistic programs, as supported by languages like Church , Venture , and Anglican . Another probabilistic programming framework is infer.NET, which implements variational message passing , an efficient algorithm for conditionally conjugate graphical models. Stan supports a more comprehensive class of models that we describe in Section 2.1.

Automatic Differentiation Variational Inference

Automatic differentiation variational inference (advi) follows a straightforward recipe. First, we transform the space of the latent variables in our model to the real coordinate space. For example, the logarithm transforms a positively constrained variable, such as a standard deviation, to the real line. Then, we posit a Gaussian variational distribution. This induces a non-Gaussian approximation in the original variable space. Last, we combine automatic differentiation with stochastic optimization to maximize the variational objective. We begin by defining the class of models we support.

Consider a dataset $\boldsymbol{\mathbf{X}}=\boldsymbol{\mathbf{x}}_{1:N}$ with $N$ observations. Each $\boldsymbol{\mathbf{x}}_{n}$ is a discrete or continuous random vector. The likelihood $p(\boldsymbol{\mathbf{X}}\mid{\theta})$ relates the observations to a set of latent random variables ${\theta}$. Bayesian analysis posits a prior density $p({\theta})$ on the latent variables. Combining the likelihood with the prior gives the joint density $p(\boldsymbol{\mathbf{X}},{\theta})=p(\boldsymbol{\mathbf{X}}\mid{\theta})\,p({\theta})$.

Many machine learning models are differentiable probability models. Linear and logistic regression, matrix factorization with continuous or discrete measurements, linear dynamical systems, and Gaussian processes are prime examples. In machine learning, we usually describe mixture models, hidden Markov models, and topic models with discrete random variables. Marginalizing out the discrete variables reveals that these are also differentiable probability models. (We show an example in Section 3.3.) Only fully discrete models, such as the Ising model, fall outside of this category.

2 Variational Inference

In Bayesian inference, we seek the posterior density $p({\theta}\mid\boldsymbol{\mathbf{X}})$, which describes how the latent variables vary, conditioned on a set of observations $\boldsymbol{\mathbf{X}}$. Many posterior densities are intractable because they lack analytic (closed-form) solutions. Thus, we seek to approximate the posterior.

Consider an approximating density $q({\theta}\,;\,{\phi})$ parameterized by ${\phi}$. We make no assumptions about its shape or support. We want to find the parameters of $q({\theta}\,;\,{\phi})$ to best match the posterior according to some loss function. Variational inference (vi) minimizes the Kullback-Leibler (kl) divergence,

from the approximation to the posterior . Typically the kl divergence also lacks an analytic form. Instead we maximize a proxy to the kl divergence, the evidence lower bound (elbo)

The first term is an expectation of the joint density under the approximation, and the second is the entropy of the variational density. Maximizing the elbo minimizes the kl divergence .

The minimization problem from Equation 1 becomes

Why is vi difficult to automate? In classical variational inference, we typically design a conditionally conjugate model; the optimal approximating family matches the prior, which satisfies the support constraint by definition . In other models, we carefully study the model and design custom approximations. These depend on the model and on the choice of the approximating density.

We adopt a transformation-based approach. First, we automatically transform the support of the latent variables in our model to the real coordinate space. Then, we posit a Gaussian variational density. The inverse of our transform induces a non-Gaussian variational approximation in the original variable space. The transformation guarantees that the non-Gaussian approximation stays within the support of the posterior. Here is how it works.

3 Automatic Transformation of Constrained Variables

and identify the transformed variables as ${\zeta}=T({\theta})$. The transformed joint density $g(\boldsymbol{\mathbf{X}},{\zeta})$ is a function of ${\zeta}$; it has the representation

where $p$ is the joint density in the original latent variable space, and $J_{T^{-1}}({\zeta})$ is the Jacobian of the inverse of $T$. Transformations of continuous probability densities require a Jacobian; it accounts for how the transformation warps unit volumes . (See Appendix D.)

As we describe in the introduction, we implement our algorithm in Stan to enable generic inference. Stan implements a model compiler that automatically handles transformations. It works by applying a library of transformations and their corresponding Jacobians to the joint model density.Stan provides transformations for upper and lower bounds, simplex and ordered vectors, and structured matrices such as covariance matrices and Cholesky factors . This transforms the joint density of any differentiable probability model to the real coordinate space. Now, we can choose a variational distribution independent from the model.

4 Implicit Non-Gaussian Variational Approximation

where the vector ${\phi}=(\mu_{1},\cdots,\mu_{K},\sigma^{2}_{1},\cdots,\sigma^{2}_{K})$ concatenates the mean and variance of each Gaussian factor. This defines our variational approximation in the real coordinate space. (Figure 3(b).)

The transformation $T$ from Equation 3 maps the support of the latent variables to the real coordinate space. Thus, its inverse $T^{-1}$ maps back to the support of the latent variables. This implicitly defines the variational approximation in the original latent variable space as \mathcal{N}(T^{-1}({\zeta})\,;\,{\mu},{\sigma}^{2})\big{|}\det J_{T^{-1}}({\zeta})\big{|}. The transformation ensures that the support of this approximation is always bounded by that of the true posterior in the original latent variable space (Figure 3(a)). Thus we can freely optimize the elbo in the real coordinate space (Figure 3(b)) without worrying about the support matching constraint.

where we plug in the analytic form for the Gaussian entropy. (Derivation in Appendix A.)

We choose a diagonal Gaussian for its efficiency and analytic entropy. This choice may call to mind the Laplace approximation technique, where a second-order Taylor expansion around the maximum-a-posteriori estimate gives a Gaussian approximation to the posterior. However, using a Gaussian variational approximation is not equivalent to the Laplace approximation . Our approach is distinct in another way: the posterior approximation in the original latent variable space (Figure 3(a)) is non-Gaussian, because of the inverse transformation $T^{-1}$ and its Jacobian.

5 Automatic Differentiation for Stochastic Optimization

We now seek to maximize the elbo in real coordinate space,

We can use gradient ascent to reach a local maximum of the elbo. Unfortunately, we cannot apply automatic differentiation to the elbo in this form. This is because the expectation defines an intractable integral that depends on ${\mu}$ and ${\sigma}^{2}$; we cannot directly represent it as a computer program. Moreover, the variance vector ${\sigma}^{2}$ must remain positive. Thus, we employ one final transformation: elliptical standardizationAlso known as a “co-ordinate transformation” , an “invertible transformation” , and the “re-parameterization trick” . , shown in Figures 3(b) and 3(c).

First, re-parameterize the Gaussian distribution with the log of the standard deviation, ${\omega}=\log({\sigma})$, applied element-wise. The support of ${\omega}$ is now the real coordinate space and ${\sigma}$ is always positive. Then, define the standardization ${\eta}=S_{{\mu},{\omega}}({\zeta})=\textrm{diag}(\exp({\omega}^{-1}))({\zeta}-{\mu})$. The standardization encapsulates the variational parameters; in return it gives a fixed variational density

The standardization transforms the variational problem from Equation 4 into

where we drop independent term from the calculation. The expectation is now in terms of the standard Gaussian, and both parameters ${\mu}$ and ${\omega}$ are unconstrained. (Figure 3(c).) We push the gradient inside the expectations and apply the chain rule to get

We can now compute the gradients inside the expectation with automatic differentiation. This leaves only the expectation. mc integration provides a simple approximation: draw $M$ samples from the standard Gaussian and evaluate the empirical mean of the gradients within the expectation . This gives unbiased noisy estimates of gradients of the elbo.

6 Scalable Automatic Variational Inference

Equipped with unbiased noisy gradients of the elbo, advi implements stochastic gradient ascent. (Algorithm 1.) We ensure convergence by choosing a decreasing step-size schedule. In practice, we use an adaptive schedule with finite memory. (See Appendix E for details.)

advi has complexity $\mathcal{O}(2NMK)$ per iteration, where $M$ is the number of mc samples (typically between 1 and 10). Coordinate ascent vi has complexity $\mathcal{O}(2NK)$ per pass over the dataset. We scale advi to large datasets using stochastic optimization . The adjustment to Algorithm 1 is simple: sample a minibatch of size $B\ll N$ from the dataset and scale the likelihood of the model by $N/B$ . The stochastic extension of advi has a per-iteration complexity $\mathcal{O}(2BMK)$.

Empirical Study

We now study advi across a variety of models. We compare its speed and accuracy to two Markov chain Monte Carlo (mcmc) sampling algorithms: Hamiltonian Monte Carlo (hmc) and the no-U-turn sampler (nuts) nuts is an adaptive extension of hmc. It is the default sampler in Stan. . We assess advi convergence by tracking the elbo; assessing convergence with mcmc techniques is less straightforward. To place advi and mcmc on a common scale, we report predictive accuracy on held-out data as a function of time. We approximate the Bayesian posterior predictive using mc integration. For the mcmc techniques, we plug in posterior samples into the likelihood. For advi, we do the same by drawing a sample from the posterior approximation at fixed intervals during the optimization. We initialize advi with a draw from a standard Gaussian.

We explore two hierarchical regression models, two matrix factorization models, and a mixture model. All of these models have nonconjugate prior structures. We conclude by analyzing a dataset of $250\,000$ images, where we report results across a range of minibatch sizes $B$.

We begin with two nonconjugate regression models: linear regression with automatic relevance determination (ard) and hierarchical logistic regression .

Linear Regression with ard. This is a sparse linear regression model with a hierarchical prior structure. (Details in Appendix F.) We simulate a dataset with $250$ regressors such that half of the regressors have no predictive power. We use $10\,000$ training samples and hold out $1000$ for testing.

Logistic Regression with Spatial Hierarchical Prior. This is a hierarchical logistic regression model from political science. The prior captures dependencies, such as states and regions, in a polling dataset from the United States 1988 presidential election . The model is nonconjugate and would require some form of approximation to derive a vi algorithm. (Details in Appendix G.)

We train using $10\,000$ data point and withhold $1536$ for evaluation. The regressors contain age, education, and state and region indicators. The dimension of the regression problem is $145$.

Results. Figure 4 plots average log predictive accuracy as a function of time. For these simple models, all methods reach the same predictive accuracy. We study advi with two settings of $M$, the number of mc samples used to estimate gradients. A single sample per iteration is sufficient; it also is the fastest. (We set $M=1$ from here on.)

2 Exploring nonconjugate Models: Non-negative Matrix Factorization

We continue by exploring two nonconjugate non-negative matrix factorization models: a constrained Gamma Poisson model and a Dirichlet Exponential model. Here, we show how easy it is to explore new models using advi. In both models, we use the Frey Face dataset, which contains $1956$ frames ($28\times 20$ pixels) of facial expressions extracted from a video sequence.

Constrained Gamma Poisson. This is a Gamma Poisson factorization model with an ordering constraint: each row of the Gamma matrix goes from small to large values. (Details in Appendix H.)

Dirichlet Exponential. This is a nonconjugate Dirichlet Exponential factorization model with a Poisson likelihood. (Details in Appendix I.)

Results. Figure 5 shows average log predictive accuracy as well as ten factors recovered from both models. advi provides an order of magnitude speed improvement over nuts (Figure 5(a)). nuts struggles with the Dirichlet Exponential model (Figure 5(b)). In both cases, hmc does not produce any useful samples within a budget of one hour; we omit hmc from the plots.

The Gamma Poisson model (Figure 5(c)) appears to pick significant frames out of the dataset. The Dirichlet Exponential factors (Figure 5(d)) are sparse and indicate components of the face that move, such as eyebrows, cheeks, and the mouth.

3 Scaling to Large Datasets: Gaussian Mixture Model

We conclude with the Gaussian mixture model (gmm) example we highlighted earlier. This is a nonconjugate gmm applied to color image histograms. We place a Dirichlet prior on the mixture proportions, a Gaussian prior on the component means, and a lognormal prior on the standard deviations. (Details in Appendix J.) We explore the imageclef dataset, which has $250\,000$ images . We withhold $10\,000$ images for evaluation.

In Figure 1(a) we randomly select $1000$ images and train a model with $10$ mixture components. nuts struggles to find an adequate solution and hmc fails altogether. This is likely due to label switching, which can affect hmc-based techniques in mixture models .

Figure 1(b) shows advi results on the full dataset. Here we use advi with stochastic subsampling of minibatches from the dataset . We increase the number of mixture components to $30$. With a minibatch size of $500$ or larger, advi reaches high predictive accuracy. Smaller minibatch sizes lead to suboptimal solutions, an effect also observed in . advi converges in about two hours.

Conclusion

We develop automatic differentiation variational inference (advi) in Stan. advi leverages automatic transformations, an implicit non-Gaussian variational approximation, and automatic differentiation. This is a valuable tool. We can explore many models, and analyze large datasets with ease. We emphasize that advi is currently available as part of Stan; it is ready for anyone to use.

We acknowledge our amazing colleagues and funding sources here.

Appendix A Transformation of the Evidence Lower Bound

Recall that ${\zeta}=T({\theta})$ and that the variational approximation in the real coordinate space is $q({\zeta}\,;\,{\mu},{\sigma}^{2})$.

We begin with the evidence lower bound (elbo) in the original latent variable space. We then transform the latent variable space of to the real coordinate space.

The variational approximation in the real coordinate space is a Gaussian. Plugging in its entropy gives the elbo in the real coordinate space

Appendix B Gradients of the Evidence Lower Bound

First, consider the gradient with respect to the ${\mu}$ parameter of the standardization. We exchange the order of the gradient and the integration through the dominated convergence theorem . The rest is the chain rule for differentiation.

Similarly, consider the gradient with respect to the ${\omega}$ parameter of the standardization. The gradient with respect to a single component, $\omega_{k}$, has a clean form. We abuse the $\nabla$ notation to maintain consistency with the rest of the text (instead of switching to $\partial$).

Appendix C Running advi in Stan

Use git to checkout the feature/bbvb branch from https://github.com/stan-dev/stan. Follow instructions to build Stan. Then download cmdStan from https://github.com/stan-dev/cmdstan. Follow instructions to build cmdStan and compile your model. You are then ready to run advi.

where myData.data.R is the dataset in the R language dump format. output_advi.csv contains samples from the posterior and elbo_advi.csv reports the elbo.

Appendix D Transformations of Continuous Probability Densities

We present a brief summary of transformations, largely based on .

Consider a univariate (scalar) random variable $X$ with probability density function $f_{X}(x)$. Let $\mathcal{X}=\textrm{supp}(f_{X}(x))$ be the support of $X$. Now consider another random variable $Y$ defined as $Y=T(X)$. Let $\mathcal{Y}=\textrm{supp}(f_{Y}(y))$ be the support of $Y$.

If $T$ is a one-to-one and differentiable function from $\mathcal{X}$ to $\mathcal{Y}$, then $Y$ has probability density function

Let us sketch a proof. Consider the cumulative density function $Y$. If the transformation $T$ is increasing, we directly apply its inverse to the cdf of $Y$. If the transformation $T$ is decreasing, we apply its inverse to one minus the cdf of $Y$. The probability density function is the derivative of the cumulative density function. These things combined give the absolute value of the derivative above.

The extension to multivariate variables $\boldsymbol{\mathbf{X}}$ and $\boldsymbol{\mathbf{Y}}$ requires a multivariate version of the absolute value of the derivative of the inverse transformation. This is the the absolute determinant of the Jacobian, $|\det J_{T^{-1}}(\boldsymbol{\mathbf{Y}})|$ where the Jacobian is

Intuitively, the Jacobian describes how a transformation warps unit volumes across spaces. This matters for transformations of random variables, since probability density functions must always integrate to one. If the transformation is linear, then we can drop the Jacobian adjustment; it evaluates to one. Similarly, affine transformations, like elliptical standardizations, do not require Jacobian adjustments; they preserve unit volumes.

Appendix E Setting a Stepsize Sequence for advi

We use adaGrad to adaptively set the stepsize sequence in advi. While adaGrad offers attractive convergence properties, in practice it can be slow because it has infinite memory. (It tracks the norm of the gradient starting from the beginning of the optimization.) In advi we randomly initialize the variational approximation, which can be far from the true posterior. This makes adaGrad take very small steps for the rest of the optimization, thus slowing convergence. Limiting adaGrad’s memory speeds up convergence in practice, an effect also observed in training neural networks . (See for an analysis of these trade-offs and a method that combines benefits from both.)

Consider the stepsize ${\rho}^{(i)}$ and a gradient vector $\boldsymbol{\mathbf{g}}^{(i)}$ at iteration $i$. The $k$th element of ${\rho}^{(i)}$ is

where, in adaGrad, $\boldsymbol{\mathbf{s}}$ is the gradient vector squared, summed over all times steps since the start of the optimization. Instead, we limit this to the past ten iterations and compute $\boldsymbol{\mathbf{s}}$ as

(In practice, we implement this recursively to save memory.) We set $\eta=0.1$ and $\tau=1$ as the default values we use in Stan.

Appendix F Linear Regression with Automatic Relevance Determination

Linear regression with automatic relevance determination (ard) is a high-dimensional sparse regression model . We describe the model below. Stan code is in Figure 6.

The inputs are $\boldsymbol{\mathbf{X}}=\boldsymbol{\mathbf{x}}_{1:N}$ where each $\boldsymbol{\mathbf{x}}_{n}$ is $D$-dimensional. The outputs are $\boldsymbol{\mathbf{y}}=y_{1:N}$ where each $y_{n}$ is $1$-dimensional. The weights vector $\boldsymbol{\mathbf{w}}$ is $D$-dimensional. The likelihood

describes measurements corrupted by iid Gaussian noise with unknown variance $\tau^{-1}$.

The ard prior and hyper-prior structure is as follows

where $\boldsymbol{\mathbf{\alpha}}$ is a $D$-dimensional hyper-prior on the weights, where each component gets its own independent Gamma prior.

We simulate data such that only half the regressions have predictive power. The results in Figure 4(a) use $a_{0}=b_{0}=c_{0}=d_{0}=1$ as hyper-parameters for the Gamma priors.

Appendix G Hierarchical Logistic Regression

Hierarchical logistic regression models dependencies in an intuitive and powerful way. We study a model of voting preferences from the 1988 United States presidential election. Chapter 14.1 of motivates the model and explains the dataset. We also describe the model below. Stan code is in Figure 7, based on .

where $\sigma(\cdot)$ is the sigmoid function (also know as the logistic function).

The variance terms all have uniform hyper-priors, constrained between 0 and 100.

Appendix H Non-negative Matrix Factorization: Constrained Gamma Poisson Model

The Gamma Poisson factorization model is a powerful way to analyze discrete data matrices .

Consider a $U\times I$ matrix of observations. We find it helpful to think of $u=\{1,\cdots,U\}$ as users and $i=\{1,\cdots,I\}$ as items, as in a recommendation system setting. The generative process for a Gamma Poisson model with $K$ factors is

For each component $k$, draw $\theta_{uk}\sim\textrm{Gam}(a_{0},b_{0})$.

For each component $k$, draw $\beta_{ik}\sim\textrm{Gam}(c_{0},d_{0})$.

Draw the observation $y_{ui}\sim\text{Poisson}({\theta}_{u}^{\top}{\beta}_{i})$.

A potential downfall of this model is that it is not uniquely identifiable: scaling ${\theta}_{u}$ by $\alpha$ and ${\beta}_{i}$ by $\alpha^{-1}$ gives the same likelihood. One way to contend with this is to constrain either vector to be a positive, ordered vector during inference. We constrain each ${\theta}_{u}$ vector in our model in this fashion. Stan code is in Figure 8. We set $K=10$ and all the Gamma hyper-parameters to 1 in our experiments.

Appendix I Non-negative Matrix Factorization: Dirichlet Exponential Model

Another model for discrete data is a Dirichlet Exponential model. The Dirichlet enforces uniqueness while the exponential promotes sparsity. This is a non-conjugate model that does not appear to have been studied in the literature.

The generative process for a Dirichlet Exponential model with $K$ factors is

Draw the $K$-vector ${\theta}_{u}\sim\text{Dir}(\boldsymbol{\mathbf{\alpha}}_{0})$.

For each component $k$, draw $\beta_{ik}\sim\text{Exponential}(\lambda_{0})$.

Draw the observation $y_{ui}\sim\text{Poisson}({\theta}_{u}^{\top}{\beta}_{i})$.

Stan code is in Figure 9. We set $K=10$, $\alpha_{0}=1000$ for each component, and $\lambda_{0}=0.1$. With this configuration of hyper-parameters, the factors ${\beta}_{i}$ are sparse and appear interpretable.

Appendix J Gaussian Mixture Model

The Gaussian mixture model (gmm) is a powerful probability model. We use it to group a dataset of natural images based on their color histograms. We build a high-dimensional gmm with a Gaussian prior for the mixture means, a lognormal prior for the mixture standard deviations, and a Dirichlet prior for the mixture components.

The images are in $\boldsymbol{\mathbf{X}}=\boldsymbol{\mathbf{x}}_{1:N}$ where each $\boldsymbol{\mathbf{x}}_{n}$ is $D$-dimensional and there are $N$ observations. The likelihood for the images is

with a Dirichlet prior for the mixture proportions

and a lognormal prior for the mixture standard deviations

The dimension of the color histograms in the imageclef dataset is $D=576$. These a concatenation of three $192$-length histograms, one for each color channel (red, green, blue) of the images.

We scale the image histograms to have zero mean and unit variance and set $\alpha_{0}=10\,000$, $\sigma_{\mu}=0.1$ and $\sigma_{\mu}$. advi code is in Figure 10. The stochastic data subsampling version of the code is in Figure 11.

References