Overdispersed Black-Box Variational Inference

Introduction

Generative probabilistic modeling is an effective approach for understanding real-world data in many areas of science (Bishop 2006; Murphy 2012). A probabilistic model describes a data-generating process through a joint distribution of observed data and latent (unobserved) variables. With a model in place, the investigator uses an inference algorithm to calculate or approximate the posterior, i.e., the conditional distribution of the latent variables given the available observations. It is through the posterior that the investigator explores the latent structure in the data and forms a predictive distribution of future data. Approximating the posterior is the central algorithmic problem for probabilistic modeling.

One of the most widely used methods to approximate the posterior distribution is variational inference (Wainwright and Jordan 2008; Jordan et al. 1999). Variational inference aims to approximate the posterior with a simpler distribution, fitting that distribution to be close to the exact posterior, where closeness is measured in terms of Kullback-Leibler (kl) divergence. In minimizing the kl, variational inference converts the problem of approximating the posterior into an optimization problem.

Traditional variational inference uses coordinate ascent to optimize its objective. This works well for models in which each conditional distribution is easy to compute (Ghahramani and Beal 2001), but is difficult to use in more complex models where the variational objective involves intractable expectations. Recent innovations in variational inference have addressed this with stochastic optimization, forming noisy gradients with Monte Carlo approximation. This strategy expands the scope of variational inference beyond traditional models, e.g., to non-conjugate probabilistic models (Carbonetto et al. 2009; Paisley et al. 2012; Salimans and Knowles 2013; Ranganath et al. 2014; Titsias and Lázaro-Gredilla 2014), deep neural networks (Neal 1992; Hinton et al. 1995; Mnih and Gregor 2014; Kingma and Welling 2014; Ranganath et al. 2015a), and probabilistic programming (Wingate and Weber 2013; Kucukelbir et al. 2015). Some of these techniques find their roots in classical policy search algorithms for reinforcement learning (Williams 1992; van de Meent et al. 2016).

These approaches must address a core problem with Monte Carlo estimates of the gradient, which is that they suffer from high variance. The estimated gradient can significantly differ from the truth and this leads to slow convergence of the optimization. There are several strategies to reduce the variance of the gradients, including Rao-Blackwellization (Casella and Robert 1996; Ranganath et al. 2014), control variates (Ross 2002; Paisley et al. 2012; Ranganath et al. 2014; Gu et al. 2016), reparameterization (Price 1958; Bonnet 1964; Salimans and Knowles 2013; Kingma and Welling 2014; Rezende et al. 2014; Kucukelbir et al. 2015), and local expectations (Titsias and Lázaro-Gredilla 2015).

In this paper we develop overdispersed black-box variational inference (o-bbvi), a new method for reducing the variance of Monte Carlo gradients in variational inference. The main idea is to use importance sampling to estimate the gradient, in order to construct a good proposal distribution that is matched to the variational problem. We show that o-bbvi applies more generally than methods such as reparameterization and local expectations, and it further improves the profile of gradients that use Rao-Blackwellization and control variates.

We demonstrate o-bbvi on two complex models: a non-conjugate time series model (Ranganath et al. 2014) and Poisson-based deep exponential families (defs) (Ranganath et al. 2015a). Our study shows that o-bbvi reduces the variance of the original black-box variational inference (bbvi) estimates (Ranganath et al. 2014), even when using only half the number of Monte Carlo samples. This provides significant savings in run-time complexity.

and form noisy estimates with samples from the proposal.

In detail, we first assume that the variational distribution is in the exponential family. (This is not an assumption about the model; most applications of variational inference use exponential family variational distributions.) We then set the proposal distribution to be in the corresponding overdispersed exponential family (Jørgensen 1987), where $\tau$ is the dispersion parameter. We show that the corresponding estimator has lower variance than the bbvi estimator, we put forward a method to adapt the dispersion parameter during optimization, and we demonstrate that this method is more efficient than bbvi. We call our approach overdispersed black-box variational inference (o-bbvi).

Organization. The rest of the paper is organized as follows. We review bbvi in Section 2. We develop o-bbvi in Section 3, describing both the basic algorithm and its extensions to adaptive proposals and high-dimensional settings. Section 4 reports on our empirical study of two non-conjugate models. We conclude the paper in Section 5.

Black-Box Variational Inference

which is a lower bound on the log of the marginal probability of the observations, $\log p(\mathbf{x})$.

With a tractable variational family (e.g., the mean-field family) and a conditionally conjugate model, A conditionally conjugate model is a model for which all the complete conditionals (i.e., the posterior distribution of each hidden variable conditioned on the observations and the rest of hidden variables) are in the same exponential family as the prior. the expectations in Eq. 5 can be computed in closed form and we can use coordinate-ascent variational inference (Ghahramani and Beal 2001). However, many models of interest are not conditionally conjugate. For these models, we need alternative methods to optimize the elbo. One approach is bbvi, which uses Monte Carlo estimates of the gradient and requires few model-specific calculations (Ranganath et al. 2014). Thus, bbvi is a variational inference algorithm that can be applied to a large class of models.

In order to reduce the variance of the estimator, bbvi uses two strategies: control variates and Rao-Blackwellization. Because we will also use these ideas in our algorithm, we briefly discuss them here.

Overdispersed Black-Box Variational Inference

We have described bbvi and its two strategies for reducing the variance of the noisy gradient. We now describe o-bbvi, a method for further reducing the variance. The main idea is to use importance sampling (Robert and Casella 2005; Rubinstein and Kroese 2011) to estimate the gradient. We first describe o-bbvi and the proposal distribution it uses. We then show that this reduces variance, discuss several important implementation details, and present the full algorithm.

where $\tau\geq 1$ is the dispersion coefficient of the overdispersed distribution (Jørgensen 1987). Hence, the o-bbvi estimator of the gradient can be expressed as

where $S$ is the number of samples of the Monte Carlo approximation.

The dispersion coefficient $\tau$ can be itself adaptive to better match the optimal proposal at each iteration of the variational optimization procedure. We put forward a method to update the value of $\tau$ in Section 3.2.

Note that our approach differs from importance weighted autoencoders (Burda et al. 2016), which also make use of importance sampling but with the goal of deriving a tighter log-likelihood lower bound in the context of the variational autoencoder (Kingma and Welling 2014). In contrast, we use importance sampling to reduce the variance of the estimator of the gradient.

After some algebra, we can express the variance of the bbvi estimator as

and we can also express the variance of the o-bbvi estimator in terms of an expectation with respect to the variational distribution as

2 Implementation

We now discuss several extensions of o-bbvi that make it more suitable for real applications.

More precisely, for the variational parameters of variable $z_{n}$, we first write the gradient as

Adaptation of the dispersion coefficients.

Our algorithm requires setting the value of the dispersion parameters $\tau_{n}$; we would like to automate this procedure. Here, we develop a method to learn these coefficients during optimization by minimizing the variance of the estimator. More precisely, we introduce stochastic gradient descent steps for $\tau_{n}$ that minimize the variance. The exact derivative of the (negative) variance with respect to $\tau_{n}$ is

where we have applied the log-derivative trick once again, as well as the extension to high dimensionality detailed above. Now a Monte Carlo estimate of this derivative can be obtained by using the same set of $S$ samples used in the update of $\lambda_{n}$. The resulting procedure is fast, with little extra overhead, since both $f_{n}(\mathbf{z})$ and $w(z_{n})$ have been pre-computed.

Thus, we perform gradient steps of the form

where $\tau_{n}$ is constrained as $\tau_{n}\geq 1$ and the derivatives are estimated via Monte Carlo approximation. Since the derivatives in Eq. 20 can be several orders of magnitude greater than $\tau_{n}$, we opt for a simple approach to choose an appropriate step size $\alpha_{n}$. In particular, we ignore the magnitude of the derivative in (20) and take a small gradient step in the direction given by its sign. Note that we do not need to satisfy the Robbins-Monro conditions here (Robbins and Monro 1951), because the adaptation of $\tau_{n}$ only defines the proposal distribution and it is not part of the original stochastic optimization procedure.

Eq. 20 can still be applied even if $\lambda_{n}$ is a vector; it only requires replacing the derivative of the variance with the summation of the derivatives for all components of $\lambda_{n}$.

Multiple importance sampling.

It may be more stable (in terms of the variance of the importance weights) to consider a set of $J$ dispersion coefficients, $\tau_{n1},\ldots,\tau_{nJ}$, instead of a single coefficient $\tau_{n}$. We propose to use a mixture with equal weights to build the proposal as follows:

where each term in the mixture is given by $r(z_{n};\lambda_{n},\tau_{nj})=g(z_{n},\tau_{nj})\exp\left\{\frac{\lambda_{n}^{\top}t(z_{n})-A(\lambda_{n})}{\tau_{nj}}\right\}$. In the importance sampling literature, this is known as multiple importance sampling (mis), as multiple proposals are used (Veach and Guibas 1995). Within the mis methods, we opt for full deterministic multiple importance sampling (dmis) because it is the approach that presents lowest variance (Hesterberg 1995; Owen and Zhou 2000; Elvira et al. 2015). In dmis, the number of samples $S$ of the Monte Carlo estimator must be an integer multiple of the number of mixture components $J$, and $S/J$ samples are deterministically assigned to each proposal $r(z_{n};\lambda_{n},\tau_{nj})$. However, the importance weights are obtained as if the samples had been actually drawn from the mixture, i.e.,

This choice of the importance weights yields an unbiased estimator with smaller variance than the standard mis approach (Owen and Zhou 2000; Elvira et al. 2015).

In the experiments in Section 4 we investigate the performance of two-component proposal distributions, where $J=2$, and compare it against our initial algorithm that uses a unique proposal, which corresponds to $J=1$. We have also conducted some additional experiments (not shown in the paper) with mixtures with higher number of components, with no significant improvements.

3 Full algorithm

We now present our full algorithm for o-bbvi. It makes use of control variates, Rao-Blackwellization, and overdispersed importance sampling with adaptation of the dispersion coefficients. At each iteration, we draw a single sample $\mathbf{z}^{(0)}$ from the variational distribution, as well as $S$ samples $z_{n}^{(s)}$ from the overdispersed proposal for each $n$ (using dmis in this step). We obtain the score function as

and the argument of the expectation in (9) as

where $p_{n}$ indicates that we use Rao-Blackwellization. Finally, the estimator of the gradient is obtained as

Following Eq. 8, we use a separate set of samples to estimate the optimal $a_{n}$ as

We use AdaGrad (Duchi et al. 2011) to obtain adaptive learning rates that ensure convergence of the stochastic optimization procedure, although other schedules can be used instead as long as they satisfy the standard Robbins-Monro conditions (Robbins and Monro 1951). In AdaGrad, the learning rate is obtained as

where $\mathbf{G}_{t}$ is a matrix that contains the sum across the first $t$ iterations of the outer products of the gradient, and $\eta$ is a constant. Thus, the stochastic gradient step is given by

where ‘$\circ$’ denotes the element-wise (Hadamard) product. Algorithm 1 summarizes the full procedure.

Empirical Study

We study our method with two non-conjugate probabilistic models: the gamma-normal time series model (gn-ts) and the Poisson deep exponential family (def). We found that overdispersed black-box variational inference (o-bbvi) reduces the variance of the black-box variational inference (bbvi) estimator and leads to faster convergence.

Models description and datasets. The gn-ts model (Ranganath et al. 2014) is a non-conjugate state-space model for sequential data that was used to showcase bbvi. The model is described by

The indices $n$, $t$, $d$ and $k$ denote observations, time instants, observation dimensions, and latent factors, respectively. The distribution GammaE denotes the expectation/variance parameterization of the gamma distribution. The model explains each datapoint $x_{ndt}$ with a latent factor model. For each time instant $t$, the mean of $x_{ndt}$ depends on the inner product $\sum_{k}z_{ntk}w_{kd}$, where $z_{ntk}$ varies smoothly across time. The variables $o_{nd}$ are an intercept that capture the baseline in the observations.

We set the hyperparameters to be $\sigma_{w}^{2}=1$, $\sigma_{o}^{2}=1$, $\sigma_{z}=1$, and $\sigma^{2}_{x}=0.01$. We use a synthetic dataset of $N=900$ time sequences of length $T=30$ and dimensionality $D=20$. We use $K=30$ latent factors, leading to $828,600$ hidden variables.

The Poisson def (Ranganath et al. 2015a) is a multi-layered latent variable model of discrete data, such as text. The model is described by

We set the prior shape and rate as $\alpha_{w}=0.1$ and $\beta_{w}=0.3$, and the prior mean for the top level of the Poisson def as $\lambda_{z}=0.1$. We use $L=3$ layers with $K=50$ latent factors each. We model the papers at the Neural Information Processing Systems (NIPS) 2011 conference. This is a data set with $D=305$ documents, $612,508$ words, and $V=5715$ vocabulary words (after removing stop words). This leads to a model with $336,500$ hidden variables.

Evaluation. We compare o-bbvi with bbvi (Ranganath et al. 2014). For a fair comparison, we use the same number of samples in both methods and estimate the inner expectation in Eq. 17 with only one sample. For the outer expectation, we use $8$ samples to estimate the gradient itself and $8$ separate samples to estimate the optimal coefficient $a_{n}$ for the control variates. For bbvi, we also doubled the number of samples to $16+16$; this is marked as “bbvi ($\times 2$)” in the plots.

For o-bbvi, we study both a single proposal and a mixture proposal with two components, respectively labeled as “o-bbvi (single proposal)” and “o-bbvi (mixture).” For the latter, we fix the dispersion coefficients $\tau_{n1}=1$ for all hidden variables and run stochastic gradient descent steps for $\tau_{n2}$.

where the held-out data contains $25\%$ randomly selected words of all documents.

Experimental setup. For each method, we initialize the variational parameters to the same point and run each algorithm with a fixed computational budget (of CPU time).

We use AdaGrad (Duchi et al. 2011) for the learning rate. We set the parameter $\eta$ in Eq. 29 to $\eta=0.5$ for the gn-ts model and $\eta=1$ for the Poisson def. When optimizing the o-bbvi dispersion coefficients $\tau_{n}$, we take steps of length $0.1$ in the direction of the (negative) gradient. We initialize the dispersion coefficients as $\tau_{n}=2$ for the single proposal and $\tau_{n2}=3$ for the two-component mixture.

We parameterize the normal distribution in terms of its mean and variance, the gamma in terms of its shape and mean, and the Poisson in terms of its mean parameter. In order to avoid constrained optimization, we apply the transformation $\lambda^{\prime}=\log(\exp(\lambda)-1)$ to those variational parameters that are constrained to be positive and take stochastic gradient steps with respect to $\lambda^{\prime}$.

Overdispersed exponential families. For a fixed dispersion coefficient $\tau$, the overdispersed exponential family of the Gaussian distribution with mean $\mu$ and variance $\sigma^{2}$ is a Gaussian distribution with mean $\mu$ and variance $\tau\sigma^{2}$. The overdispersed gamma distribution with shape $s$ and rate $r$ is given by a new gamma distribution with shape $\frac{s+\tau-1}{\tau}$ and rate $\frac{r}{\tau}$. The overdispersed $\textrm{Poisson}(\lambda)$ distribution is a $\textrm{Poisson}(\lambda^{1/\tau})$ distribution.

2 Results

Figures 1 and 2 show the evolution of the elbo, the predictive performance, and the average sample variance of the estimator for both models and all methods. We plot these metrics as a function of running time, and each method is run with the same computational budget.

For the gn-ts model, Figure 1a shows that the variance of o-bbvi is significantly lower than bbvi and bbvi with twice the number of samples. Additionally, Figures 1b and 1c show that o-bbvi outperforms vanilla bbvi in terms of both elbo and held-out likelihood. According to these figures, using a single or mixture proposal does not seem to significantly affect performance.

The results on the Poisson def are similar (Figure 2). Figure 2a shows the average sample variance of the estimator; again, o-bbvi outperforms both bbvi algorithms. Figures 2b and 2c show the evolution of the elbo and the held-out perplexity, respectively, where o-bbvi also outperforms bbvi. Here, the two-component mixture proposal performs slightly better than the single proposal. This is consistent with Figure 2a, which indicates that the mixture proposal gives more stable estimates than the single proposal.

Finally, for the gn-ts model only, we also apply the local expectations algorithm of Titsias and Lázaro-Gredilla 2015, which relies on exact or numerical integration to reduce the variance of the estimator. We form noisy gradients using numerical quadratures for the Gaussian random variables and standard bbvi for the gamma variables (these results are not plotted in the paper). We found that local expectations accurately approximate the gradient for the Gaussian distributions. It converges slightly faster at the beginning of the run, although o-bbvi quickly reaches the same performance. (We conjecture that it is faster because the local expectations algorithm does not require the use of control variates. This saves evaluations of the log-joint probability of the model and thus it can run more iterations in the same period of time.)

However, we emphasize that o-bbvi is a more general algorithm than local expectations. The local expectations of Titsias and Lázaro-Gredilla 2015 are only available for discrete distributions with finite support and for continuous distributions for which numerical quadratures are accurate (such as Gaussian distributions). They fail to approximate the expectations for other exponential family distributions (e.g., gamma, Although the univariate gamma distribution is amenable to numerical integration, we have found that the approximation of the expectations are not accurate when the shape parameter of the gamma distribution is below $1$, due to the singularity at $0$. Poisson, and others). For example, they cannot handle the Poisson def.

Conclusions

We have developed overdispersed black-box variational inference (o-bbvi), a method that relies on importance sampling to reduce the variance of the stochastic gradients in black-box variational inference (bbvi). o-bbvi uses an importance sampling proposal distribution that has heavier tails than the actual variational distribution. In particular, we choose the proposal as an overdispersed distribution in the same exponential family as the variational distribution. Like bbvi, our approach is amenable to mean field or structured variational inference, as well as variational models (Ranganath et al. 2015b; Tran et al. 2016).

We have studied the performance of our method on two complex probabilistic models. Our results show that bbvi effectively benefits from the use of overdispersed importance sampling, and o-bbvi leads to faster convergence in the resulting stochastic optimization procedure.

There are several avenues for future work. First, we can explore other proposal distributions to provide a better fit to the optimal ones while still maintaining computational efficiency. Second, we can apply quasi-Monte Carlo methods to further decrease the sampling variance, as already suggested by Ranganath et al. 2014. Finally, we can combine the reparameterization trick with overdispersed proposals to explore whether variance is further reduced.

References